[Enchance] lazy import for actions (#146)

lagent/actions/arxiv_search.py

@@ -1,7 +1,5 @@
 from typing import Optional, Type
 
-import arxiv
-
 from lagent.actions.base_action import BaseAction, tool_api
 from lagent.actions.parser import BaseParser, JsonParser
 from lagent.schema import ActionReturn, ActionStatusCode
@@ -37,6 +35,8 @@ def get_arxiv_article_information(self, query: str) -> dict:
             :class:`dict`: article information
                 * content (str): a list of 3 arxiv search papers
         """
+        import arxiv
+
         try:
             results = arxiv.Search(  # type: ignore
                 query[:self.max_query_len],

lagent/actions/google_scholar_search.py

@@ -2,8 +2,6 @@
 import os
 from typing import Optional, Type
 
-from serpapi import GoogleSearch
-
 from lagent.actions.base_action import BaseAction, tool_api
 from lagent.schema import ActionReturn, ActionStatusCode
 from .parser import BaseParser, JsonParser
@@ -78,6 +76,7 @@ def search_google_scholar(
                 - organic_id: a list of the organic results' ids of the three selected papers
                 - pub_info: publication information of selected papers
         """
+        from serpapi import GoogleSearch
         params = {
             'q': query,
             'engine': 'google_scholar',
@@ -154,6 +153,7 @@ def get_author_information(self,
                 * articles: at most 3 articles by the author
                 * website: the author's homepage url
         """
+        from serpapi import GoogleSearch
         params = {
             'engine': 'google_scholar_author',
             'author_id': author_id,
@@ -204,6 +204,7 @@ def get_citation_format(self,
                 * authors: the authors of the article
                 * citation: the citation format of the article
         """
+        from serpapi import GoogleSearch
         params = {
             'q': q,
             'engine': 'google_scholar_cite',
@@ -246,6 +247,7 @@ def get_author_id(self,
             :class:`dict`: author id
                 * author_id: the author_id of the author
         """
+        from serpapi import GoogleSearch
         params = {
             'mauthors': mauthors,
             'engine': 'google_scholar_profiles',

lagent/actions/ipython_interactive.py

@@ -5,10 +5,6 @@
 from io import StringIO
 from typing import Optional, Type
 
-import json5
-from IPython import InteractiveShell
-from timeout_decorator import timeout as timer
-
 from ..schema import ActionReturn, ActionStatusCode
 from .base_action import BaseAction, tool_api
 from .parser import BaseParser, JsonParser
@@ -55,6 +51,7 @@ def __init__(
         enable: bool = True,
     ):
         super().__init__(description, parser, enable)
+        from IPython import InteractiveShell
         self.timeout = timeout
         self._executor = InteractiveShell()
         self._highlighting = re.compile(r'\x1b\[\d{,3}(;\d{,3}){,3}m')
@@ -74,6 +71,7 @@ def run(self, command: str, timeout: Optional[int] = None) -> ActionReturn:
             timeout (:class:`Optional[int]`): timeout for execution.
                 This argument only works in the main thread. Defaults to ``None``.
         """
+        from timeout_decorator import timeout as timer
         tool_return = ActionReturn(args={'text': command}, type=self.name)
         ret = (
             timer(timeout or self.timeout)(self.exec)(command)
@@ -171,6 +169,8 @@ def extract_code(text: str) -> str:
         Returns:
             :class:`str`: Python code
         """
+        import json5
+
         # Match triple backtick blocks first
         triple_match = re.search(r'```[^\n]*\n(.+?)```', text, re.DOTALL)
         # Match single backtick blocks second

lagent/actions/ipython_interpreter.py

@@ -11,10 +11,6 @@
 import uuid
 from typing import Optional, Tuple, Type
 
-import json5
-import PIL.Image
-from jupyter_client import KernelManager
-
 from lagent.actions.base_action import BaseAction, tool_api
 from lagent.actions.parser import BaseParser, JsonParser
 from lagent.schema import ActionReturn, ActionStatusCode
@@ -75,6 +71,8 @@ def __init__(self,
 
     @staticmethod
     def start_kernel():
+        from jupyter_client import KernelManager
+
         # start the kernel and manager
         km = KernelManager()
         km.start_kernel()
@@ -235,6 +233,8 @@ def run(self, command: str, timeout: Optional[int] = None) -> ActionReturn:
 
 
 def extract_code(text):
+    import json5
+
     # Match triple backtick blocks first
     triple_match = re.search(r'```[^\n]*\n(.+?)```', text, re.DOTALL)
     # Match single backtick blocks second
@@ -258,6 +258,7 @@ def escape_ansi(line):
 
 
 def publish_image_to_local(image_base64: str, work_dir='./work_dir/tmp_dir'):
+    import PIL.Image
     image_file = str(uuid.uuid4()) + '.png'
     local_image_file = os.path.join(work_dir, image_file)
 

lagent/actions/ppt.py

@@ -1,7 +1,5 @@
 from typing import Dict, Optional, Type
 
-from pptx import Presentation
-
 from lagent.actions.base_action import BaseAction, tool_api
 from lagent.actions.parser import BaseParser, JsonParser
 
@@ -40,6 +38,7 @@ def create_file(self, theme: str, abs_location: str) -> dict:
             :class:`dict`: operation status
                 * status: the result of the execution
         """
+        from pptx import Presentation
         self.location = abs_location
         try:
             self.pointer = Presentation(self.theme_mapping[theme]['template'])

lagent/actions/python_interpreter.py

@@ -4,8 +4,6 @@
 from contextlib import redirect_stdout
 from typing import Any, Optional, Type
 
-from func_timeout import FunctionTimedOut, func_set_timeout
-
 from lagent.actions.base_action import BaseAction, tool_api
 from lagent.actions.parser import BaseParser, JsonParser
 from lagent.schema import ActionReturn, ActionStatusCode
@@ -85,6 +83,7 @@ def solution():
         Args:
             command (:class:`str`): Python code snippet
         """
+        from func_timeout import FunctionTimedOut, func_set_timeout
         self.runtime = GenericRuntime()
         try:
             tool_return = func_set_timeout(self.timeout)(self._call)(command)

requirements/optional.txt

@@ -1,4 +1,7 @@
+google-search-results
 lmdeploy>=0.2.2
-streamlit
+pillow
+python-pptx
+timeout_decorator
 torch
 transformers>=4.34

requirements/runtime.txt

@@ -1,16 +1,13 @@
 arxiv
 distro
 func_timeout
-google-search-results
 griffe
 json5
 jsonschema
 jupyter
 jupyter_client
 phx-class-registry
-pillow
-python-pptx
 requests
+streamlit
 tiktoken
-timeout_decorator
 typing-extensions