mirror of
https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git
synced 2025-07-06 02:25:10 +08:00
94 lines
3.5 KiB
Python
94 lines
3.5 KiB
Python
import logging
|
|
from threading import Lock
|
|
from typing import Union
|
|
|
|
import contexts
|
|
from core.datasource.__base.datasource_plugin import DatasourcePlugin
|
|
from core.datasource.__base.datasource_provider import DatasourcePluginProviderController
|
|
from core.datasource.entities.common_entities import I18nObject
|
|
from core.datasource.entities.datasource_entities import DatasourceProviderType
|
|
from core.datasource.errors import DatasourceProviderNotFoundError
|
|
from core.plugin.impl.tool import PluginToolManager
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class DatasourceManager:
|
|
_builtin_provider_lock = Lock()
|
|
_hardcoded_providers: dict[str, DatasourcePluginProviderController] = {}
|
|
_builtin_providers_loaded = False
|
|
_builtin_tools_labels: dict[str, Union[I18nObject, None]] = {}
|
|
|
|
@classmethod
|
|
def get_datasource_plugin_provider(cls, provider: str, tenant_id: str) -> DatasourcePluginProviderController:
|
|
"""
|
|
get the datasource plugin provider
|
|
"""
|
|
# check if context is set
|
|
try:
|
|
contexts.datasource_plugin_providers.get()
|
|
except LookupError:
|
|
contexts.datasource_plugin_providers.set({})
|
|
contexts.datasource_plugin_providers_lock.set(Lock())
|
|
|
|
with contexts.datasource_plugin_providers_lock.get():
|
|
datasource_plugin_providers = contexts.datasource_plugin_providers.get()
|
|
if provider in datasource_plugin_providers:
|
|
return datasource_plugin_providers[provider]
|
|
|
|
manager = PluginToolManager()
|
|
provider_entity = manager.fetch_datasource_provider(tenant_id, provider)
|
|
if not provider_entity:
|
|
raise DatasourceProviderNotFoundError(f"plugin provider {provider} not found")
|
|
|
|
controller = DatasourcePluginProviderController(
|
|
entity=provider_entity.declaration,
|
|
plugin_id=provider_entity.plugin_id,
|
|
plugin_unique_identifier=provider_entity.plugin_unique_identifier,
|
|
tenant_id=tenant_id,
|
|
)
|
|
|
|
datasource_plugin_providers[provider] = controller
|
|
|
|
return controller
|
|
|
|
@classmethod
|
|
def get_datasource_runtime(
|
|
cls,
|
|
provider_type: DatasourceProviderType,
|
|
provider_id: str,
|
|
datasource_name: str,
|
|
tenant_id: str,
|
|
) -> DatasourcePlugin:
|
|
"""
|
|
get the datasource runtime
|
|
|
|
:param provider_type: the type of the provider
|
|
:param provider_id: the id of the provider
|
|
:param datasource_name: the name of the datasource
|
|
:param tenant_id: the tenant id
|
|
|
|
:return: the datasource plugin
|
|
"""
|
|
if provider_type == DatasourceProviderType.RAG_PIPELINE:
|
|
return cls.get_datasource_plugin_provider(provider_id, tenant_id).get_datasource(datasource_name)
|
|
else:
|
|
raise DatasourceProviderNotFoundError(f"provider type {provider_type.value} not found")
|
|
|
|
@classmethod
|
|
def list_datasource_providers(cls, tenant_id: str) -> list[DatasourcePluginProviderController]:
|
|
"""
|
|
list all the datasource providers
|
|
"""
|
|
manager = PluginToolManager()
|
|
provider_entities = manager.fetch_datasources(tenant_id)
|
|
return [
|
|
DatasourcePluginProviderController(
|
|
entity=provider.declaration,
|
|
plugin_id=provider.plugin_id,
|
|
plugin_unique_identifier=provider.plugin_unique_identifier,
|
|
tenant_id=tenant_id,
|
|
)
|
|
for provider in provider_entities
|
|
]
|