mirror of
https://github.com/onyx-dot-app/onyx.git
synced 2026-10-23 11:53:08 -05:00
58 lines
1.3 KiB
Python
58 lines
1.3 KiB
Python
import functools
|
|
import importlib
|
|
from typing import Any
|
|
from typing import TypeVar
|
|
|
|
from danswer.utils.logger import setup_logger
|
|
|
|
|
|
logger = setup_logger()
|
|
|
|
|
|
class DanswerVersion:
|
|
def __init__(self) -> None:
|
|
self._is_ee = False
|
|
|
|
def set_ee(self) -> None:
|
|
self._is_ee = True
|
|
|
|
def get_is_ee_version(self) -> bool:
|
|
return self._is_ee
|
|
|
|
|
|
global_version = DanswerVersion()
|
|
|
|
|
|
@functools.lru_cache(maxsize=128)
|
|
def fetch_versioned_implementation(module: str, attribute: str) -> Any:
|
|
logger.debug("Fetching versioned implementation for %s.%s", module, attribute)
|
|
is_ee = global_version.get_is_ee_version()
|
|
|
|
module_full = f"ee.{module}" if is_ee else module
|
|
try:
|
|
return getattr(importlib.import_module(module_full), attribute)
|
|
except ModuleNotFoundError:
|
|
# try the non-ee version as a fallback
|
|
if is_ee:
|
|
return getattr(importlib.import_module(module), attribute)
|
|
|
|
raise
|
|
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
def fetch_versioned_implementation_with_fallback(
|
|
module: str, attribute: str, fallback: T
|
|
) -> T:
|
|
try:
|
|
return fetch_versioned_implementation(module, attribute)
|
|
except Exception as e:
|
|
logger.warning(
|
|
"Failed to fetch versioned implementation for %s.%s: %s",
|
|
module,
|
|
attribute,
|
|
e,
|
|
)
|
|
return fallback
|