Delay LiteLLM import to improve speed on decompiler init

mahaloz · Aug 26, 2024 · 86d0db2 · 86d0db2
1 parent c8ddbd7
commit 86d0db2
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/dailalib/__init__.py b/dailalib/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "3.4.4"
+__version__ = "3.4.5"
 
 from .api import AIAPI, LiteLLMAIAPI
 from libbs.api import DecompilerInterface

diff --git a/dailalib/api/litellm/litellm_api.py b/dailalib/api/litellm/litellm_api.py
@@ -2,7 +2,6 @@
 from typing import Optional
 import os
 
-from litellm import completion
 import tiktoken
 
 from ..ai_api import AIAPI
@@ -52,6 +51,9 @@ def query_model(
         model: Optional[str] = None,
         max_tokens=None,
     ):
+        # delay import because litellm attempts to query the server on import to collect cost information.
+        from litellm import completion
+
         if not self.api_key:
             raise ValueError(f"Model API key is not set. Please set it before querying the model {self.model}")