LiteLLM Lazy loading

jac-byllm/byllm/llm.py

@@ -17,7 +17,6 @@
 # https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json
 os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
 
-from .llm_connector import LLMConnector
 from .types import CompletionResult
 
 SYSTEM_PERSONA = """\
@@ -48,6 +47,8 @@ def __init__(self, model_name: str, **kwargs: object) -> None:
             api_key: API key for the model provider
             **kwargs: Additional configuration options
         """
+        from .llm_connector import LLMConnector
+
         self.llm_connector = LLMConnector.for_model(model_name, **kwargs)
 
     def __call__(self, **kwargs: object) -> "Model":

jac-byllm/byllm/llm_connector.py

@@ -22,14 +22,14 @@
 
 import litellm
 from litellm._logging import _disable_debugging
+from litellm.types.utils import Message as LiteLLMMessage
 
 from openai import OpenAI
 
 from .mtir import MTIR
 
 from .types import (
     CompletionResult,
-    LiteLLMMessage,
     MockToolCall,
     ToolCall,
 )

jac-byllm/byllm/mtir.py

@@ -8,7 +8,6 @@
 
 from byllm.schema import json_to_instance, type_to_schema
 from byllm.types import (
-    LiteLLMMessage,
     Media,
     Message,
     MessageRole,
@@ -150,7 +149,7 @@ def add_message(self, message: MessageType) -> None:
         """Add a message to the request."""
         self.messages.append(message)
 
-    def get_msg_list(self) -> list[dict[str, object] | LiteLLMMessage]:
+    def get_msg_list(self) -> "list[dict[str, object] | MessageType]":
         """Return the messages in a format suitable for LLM API."""
         return [
             msg.to_dict() if isinstance(msg, Message) else msg for msg in self.messages

jac-byllm/byllm/types.py

@@ -12,19 +12,22 @@
 from dataclasses import dataclass
 from enum import StrEnum
 from io import BytesIO
-from typing import Callable, TypeAlias, get_type_hints
+from typing import Callable, TYPE_CHECKING, TypeAlias, get_type_hints
 
 from PIL.Image import open as open_image
 
-from litellm.types.utils import Message as LiteLLMMessage
-
 from pydantic import TypeAdapter
 
 from .schema import tool_to_schema
 
-# The message can be a jaclang defined message or what ever the llm
-# returned object that was feed back to the llm as it was given (dict).
-MessageType: TypeAlias = "Message | LiteLLMMessage"
+if TYPE_CHECKING:
+    from litellm.types.utils import Message as LiteLLMMessage
+
+    # The message can be a jaclang defined message or what ever the llm
+    # returned object that was feed back to the llm as it was given (dict).
+    MessageType: TypeAlias = "Message | LiteLLMMessage"
+else:
+    MessageType: TypeAlias = "Message"
 
 
 class MessageRole(StrEnum):