From 86d0db2c4d146cca9941484f41b6bff455017003 Mon Sep 17 00:00:00 2001
From: mahaloz <zion@zionbasque.com>
Date: Mon, 26 Aug 2024 09:14:19 -0700
Subject: [PATCH] Delay LiteLLM import to improve speed on decompiler init

---
 dailalib/__init__.py                | 2 +-
 dailalib/api/litellm/litellm_api.py | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/dailalib/__init__.py b/dailalib/__init__.py
index 2b3f550..89c0a1f 100644
--- 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
index 778467e..ead738b 100644
--- 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}")