Move 'timeout' to common

JinHai-CN · JinHai-CN · commit c17214dac982 · 2025-11-04T11:12:08.000+08:00
Signed-off-by: Jin Hai &lt;haijin.chn@gmail.com&gt;
diff --git a/agent/component/agent_with_tools.py b/agent/component/agent_with_tools.py
@@ -27,7 +27,7 @@
 from api.db.services.llm_service import LLMBundle
 from api.db.services.tenant_llm_service import TenantLLMService
 from api.db.services.mcp_server_service import MCPServerService
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 from rag.prompts.generator import next_step, COMPLETE_TASK, analyze_task, \
     citation_prompt, reflect, rank_memories, kb_prompt, citation_plus, full_question, message_fit_in
 from rag.utils.mcp_tool_call_conn import MCPToolCallSession, mcp_tool_metadata_to_openai_tool
diff --git a/agent/component/base.py b/agent/component/base.py
@@ -25,7 +25,7 @@
 import pandas as pd
 import trio
 from agent import settings
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 
 
 _FEEDED_DEPRECATED_PARAMS = "_feeded_deprecated_params"
diff --git a/agent/component/categorize.py b/agent/component/categorize.py
@@ -21,7 +21,7 @@
 from api.db import LLMType
 from api.db.services.llm_service import LLMBundle
 from agent.component.llm import LLMParam, LLM
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 from rag.llm.chat_model import ERROR_PREFIX
 
 
diff --git a/agent/component/invoke.py b/agent/component/invoke.py
@@ -23,7 +23,7 @@
 import requests
 
 from agent.component.base import ComponentBase, ComponentParamBase
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 from deepdoc.parser import HtmlParser
 
 
diff --git a/agent/component/llm.py b/agent/component/llm.py
@@ -25,7 +25,7 @@
 from api.db.services.llm_service import LLMBundle
 from api.db.services.tenant_llm_service import TenantLLMService
 from agent.component.base import ComponentBase, ComponentParamBase
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 from rag.prompts.generator import tool_call_summary, message_fit_in, citation_prompt, structured_output_prompt
 
 
diff --git a/agent/component/message.py b/agent/component/message.py
@@ -23,7 +23,7 @@
 from agent.component.base import ComponentBase, ComponentParamBase
 from jinja2 import Template as Jinja2Template
 
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 
 
 class MessageParam(ComponentParamBase):
diff --git a/agent/component/string_transform.py b/agent/component/string_transform.py
@@ -18,7 +18,7 @@
 from abc import ABC
 from jinja2 import Template as Jinja2Template
 from agent.component.base import ComponentParamBase
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 from .message import Message
 
 
diff --git a/agent/component/switch.py b/agent/component/switch.py
@@ -19,7 +19,7 @@
 from typing import Any
 
 from agent.component.base import ComponentBase, ComponentParamBase
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 
 
 class SwitchParam(ComponentParamBase):
diff --git a/agent/tools/arxiv.py b/agent/tools/arxiv.py
@@ -19,7 +19,7 @@
 from abc import ABC
 import arxiv
 from agent.tools.base import ToolParamBase, ToolMeta, ToolBase
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 
 
 class ArXivParam(ToolParamBase):
diff --git a/agent/tools/code_exec.py b/agent/tools/code_exec.py
@@ -22,7 +22,7 @@
 from pydantic import BaseModel, Field, field_validator
 from agent.tools.base import ToolParamBase, ToolBase, ToolMeta
 from api import settings
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 
 
 class Language(StrEnum):
diff --git a/agent/tools/duckduckgo.py b/agent/tools/duckduckgo.py
@@ -19,7 +19,7 @@
 from abc import ABC
 from duckduckgo_search import DDGS
 from agent.tools.base import ToolMeta, ToolParamBase, ToolBase
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 
 
 class DuckDuckGoParam(ToolParamBase):
diff --git a/agent/tools/email.py b/agent/tools/email.py
@@ -25,7 +25,7 @@
 from email.utils import formataddr
 
 from agent.tools.base import ToolParamBase, ToolBase, ToolMeta
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 
 
 class EmailParam(ToolParamBase):
diff --git a/agent/tools/exesql.py b/agent/tools/exesql.py
@@ -22,7 +22,7 @@
 import psycopg2
 import pyodbc
 from agent.tools.base import ToolParamBase, ToolBase, ToolMeta
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 
 
 class ExeSQLParam(ToolParamBase):
diff --git a/agent/tools/github.py b/agent/tools/github.py
@@ -19,7 +19,7 @@
 from abc import ABC
 import requests
 from agent.tools.base import ToolParamBase, ToolMeta, ToolBase
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 
 
 class GitHubParam(ToolParamBase):
diff --git a/agent/tools/google.py b/agent/tools/google.py
@@ -19,7 +19,7 @@
 from abc import ABC
 from serpapi import GoogleSearch
 from agent.tools.base import ToolParamBase, ToolMeta, ToolBase
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 
 
 class GoogleParam(ToolParamBase):
diff --git a/agent/tools/googlescholar.py b/agent/tools/googlescholar.py
@@ -19,7 +19,7 @@
 from abc import ABC
 from scholarly import scholarly
 from agent.tools.base import ToolMeta, ToolParamBase, ToolBase
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 
 
 class GoogleScholarParam(ToolParamBase):
diff --git a/agent/tools/pubmed.py b/agent/tools/pubmed.py
@@ -21,7 +21,7 @@
 import re
 import xml.etree.ElementTree as ET
 from agent.tools.base import ToolParamBase, ToolMeta, ToolBase
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 
 
 class PubMedParam(ToolParamBase):
diff --git a/agent/tools/retrieval.py b/agent/tools/retrieval.py
@@ -25,7 +25,7 @@
 from api.db.services.knowledgebase_service import KnowledgebaseService
 from api.db.services.llm_service import LLMBundle
 from api import settings
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 from rag.app.tag import label_question
 from rag.prompts.generator import cross_languages, kb_prompt, gen_meta_filter
 
diff --git a/agent/tools/searxng.py b/agent/tools/searxng.py
@@ -19,7 +19,7 @@
 from abc import ABC
 import requests
 from agent.tools.base import ToolMeta, ToolParamBase, ToolBase
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 
 
 class SearXNGParam(ToolParamBase):
diff --git a/agent/tools/tavily.py b/agent/tools/tavily.py
@@ -19,7 +19,7 @@
 from abc import ABC
 from tavily import TavilyClient
 from agent.tools.base import ToolParamBase, ToolBase, ToolMeta
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 
 
 class TavilySearchParam(ToolParamBase):
diff --git a/agent/tools/wencai.py b/agent/tools/wencai.py
@@ -21,7 +21,7 @@
 import pywencai
 
 from agent.tools.base import ToolParamBase, ToolMeta, ToolBase
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 
 
 class WenCaiParam(ToolParamBase):
diff --git a/agent/tools/wikipedia.py b/agent/tools/wikipedia.py
@@ -19,7 +19,7 @@
 from abc import ABC
 import wikipedia
 from agent.tools.base import ToolMeta, ToolParamBase, ToolBase
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 
 
 class WikipediaParam(ToolParamBase):
diff --git a/agent/tools/yahoofinance.py b/agent/tools/yahoofinance.py
@@ -20,7 +20,7 @@
 import pandas as pd
 import yfinance as yf
 from agent.tools.base import ToolMeta, ToolParamBase, ToolBase
-from api.utils.api_utils import timeout
+from common.connection_utils import timeout
 
 
 class YahooFinanceParam(ToolParamBase):
diff --git a/api/utils/api_utils.py b/api/utils/api_utils.py
@@ -13,17 +13,14 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
-import asyncio
+
 import functools
 import json
 import logging
 import os
-import queue
-import threading
 import time
 from copy import deepcopy
 from functools import wraps
-from typing import Any, Callable, Coroutine, Optional, Type, Union
 
 import requests
 import trio
@@ -43,6 +40,7 @@
 from api.db.db_models import APIToken
 from api.utils.json_encode import CustomJSONEncoder
 from rag.utils.mcp_tool_call_conn import MCPToolCallSession, close_multiple_mcp_toolcall_sessions
+from common.connection_utils import timeout
 
 requests.models.complexjson.dumps = functools.partial(json.dumps, cls=CustomJSONEncoder)
 
@@ -604,82 +602,6 @@ def get_mcp_tools(mcp_servers: list, timeout: float | int = 10) -> tuple[dict, s
         return {}, str(e)
 
 
-TimeoutException = Union[Type[BaseException], BaseException]
-OnTimeoutCallback = Union[Callable[..., Any], Coroutine[Any, Any, Any]]
-
-
-def timeout(seconds: float | int | str = None, attempts: int = 2, *, exception: Optional[TimeoutException] = None, on_timeout: Optional[OnTimeoutCallback] = None):
-    if isinstance(seconds, str):
-        seconds = float(seconds)
-    def decorator(func):
-        @wraps(func)
-        def wrapper(*args, **kwargs):
-            result_queue = queue.Queue(maxsize=1)
-
-            def target():
-                try:
-                    result = func(*args, **kwargs)
-                    result_queue.put(result)
-                except Exception as e:
-                    result_queue.put(e)
-
-            thread = threading.Thread(target=target)
-            thread.daemon = True
-            thread.start()
-
-            for a in range(attempts):
-                try:
-                    if os.environ.get("ENABLE_TIMEOUT_ASSERTION"):
-                        result = result_queue.get(timeout=seconds)
-                    else:
-                        result = result_queue.get()
-                    if isinstance(result, Exception):
-                        raise result
-                    return result
-                except queue.Empty:
-                    pass
-            raise TimeoutError(f"Function '{func.__name__}' timed out after {seconds} seconds and {attempts} attempts.")
-
-        @wraps(func)
-        async def async_wrapper(*args, **kwargs) -> Any:
-            if seconds is None:
-                return await func(*args, **kwargs)
-
-            for a in range(attempts):
-                try:
-                    if os.environ.get("ENABLE_TIMEOUT_ASSERTION"):
-                        with trio.fail_after(seconds):
-                            return await func(*args, **kwargs)
-                    else:
-                        return await func(*args, **kwargs)
-                except trio.TooSlowError:
-                    if a < attempts - 1:
-                        continue
-                    if on_timeout is not None:
-                        if callable(on_timeout):
-                            result = on_timeout()
-                            if isinstance(result, Coroutine):
-                                return await result
-                            return result
-                        return on_timeout
-
-                    if exception is None:
-                        raise TimeoutError(f"Operation timed out after {seconds} seconds and {attempts} attempts.")
-
-                    if isinstance(exception, BaseException):
-                        raise exception
-
-                    if isinstance(exception, type) and issubclass(exception, BaseException):
-                        raise exception(f"Operation timed out after {seconds} seconds and {attempts} attempts.")
-
-                    raise RuntimeError("Invalid exception type provided")
-
-        if asyncio.iscoroutinefunction(func):
-            return async_wrapper
-        return wrapper
-
-    return decorator
-
 
 async def is_strong_enough(chat_model, embedding_model):
     count = settings.STRONG_TEST_COUNT
diff --git a/common/connection_utils.py b/common/connection_utils.py
diff --git a/deepdoc/parser/figure_parser.py b/deepdoc/parser/figure_parser.py
diff --git a/graphrag/general/community_reports_extractor.py b/graphrag/general/community_reports_extractor.py
diff --git a/graphrag/general/extractor.py b/graphrag/general/extractor.py
diff --git a/graphrag/general/index.py b/graphrag/general/index.py
diff --git a/graphrag/utils.py b/graphrag/utils.py
diff --git a/rag/flow/base.py b/rag/flow/base.py
diff --git a/rag/flow/tokenizer/tokenizer.py b/rag/flow/tokenizer/tokenizer.py
diff --git a/rag/raptor.py b/rag/raptor.py
diff --git a/rag/svr/task_executor.py b/rag/svr/task_executor.py
