mirror of
https://github.com/langgenius/dify.git
synced 2026-08-24 12:32:54 -04:00
Stephen Zhou
220 lines
10 KiB
Python
220 lines
10 KiB
Python
"""Read-only datasource discovery used before a KnowledgeFS Space exists."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from collections.abc import Callable
|
|
from typing import Any, cast
|
|
|
|
from sqlalchemy import select
|
|
|
|
from core.datasource.datasource_manager import DatasourceManager
|
|
from core.datasource.entities.datasource_entities import (
|
|
DatasourceProviderType,
|
|
OnlineDriveBrowseFilesRequest,
|
|
)
|
|
from core.datasource.online_document.online_document_plugin import OnlineDocumentDatasourcePlugin
|
|
from core.datasource.online_drive.online_drive_plugin import OnlineDriveDatasourcePlugin
|
|
from core.datasource.website_crawl.website_crawl_plugin import WebsiteCrawlDatasourcePlugin
|
|
from models.account import Account
|
|
from models.credential_permission import CredentialType
|
|
from models.oauth import DatasourceProvider
|
|
from services.credential_permission_service import CredentialPermissionService
|
|
from services.datasource_provider_service import DatasourceProviderService
|
|
from services.knowledge_fs.product_dto import (
|
|
KnowledgeFSInitialSourcePreviewDocumentResponse,
|
|
KnowledgeFSInitialSourcePreviewFileResponse,
|
|
KnowledgeFSInitialSourcePreviewPageResponse,
|
|
KnowledgeFSInitialSourcePreviewPayload,
|
|
KnowledgeFSInitialSourcePreviewResponse,
|
|
KnowledgeFSInitialWebsiteSourcePreviewPayload,
|
|
)
|
|
|
|
_MAX_PREVIEW_ITEMS = 200
|
|
|
|
|
|
class KnowledgeFSInitialSourcePreviewCanceledError(RuntimeError):
|
|
pass
|
|
|
|
|
|
def _raise_if_canceled(is_canceled: Callable[[], bool] | None) -> None:
|
|
if is_canceled is not None and is_canceled():
|
|
raise KnowledgeFSInitialSourcePreviewCanceledError("Datasource preview was canceled")
|
|
|
|
|
|
class KnowledgeFSInitialSourcePreviewService:
|
|
def __init__(self, session_maker) -> None:
|
|
self._session_maker = session_maker
|
|
|
|
def require_visible_credential(
|
|
self,
|
|
*,
|
|
tenant_id: str,
|
|
account: Account,
|
|
payload: KnowledgeFSInitialSourcePreviewPayload | KnowledgeFSInitialWebsiteSourcePreviewPayload,
|
|
) -> None:
|
|
query = select(DatasourceProvider).where(
|
|
DatasourceProvider.tenant_id == tenant_id,
|
|
DatasourceProvider.id == payload.credential_id,
|
|
DatasourceProvider.provider == payload.provider,
|
|
DatasourceProvider.plugin_id == payload.plugin_id,
|
|
)
|
|
query = CredentialPermissionService.apply_visibility_filter(
|
|
query,
|
|
model_id_column=DatasourceProvider.id,
|
|
model_user_id_column=DatasourceProvider.user_id,
|
|
model_visibility_column=DatasourceProvider.visibility,
|
|
credential_type=CredentialType.DATASOURCE_PROVIDER,
|
|
user=account,
|
|
)
|
|
with self._session_maker() as session:
|
|
if session.scalar(query.limit(1)) is None:
|
|
raise PermissionError("Datasource credential is unavailable")
|
|
|
|
def preview(
|
|
self,
|
|
*,
|
|
tenant_id: str,
|
|
account: Account,
|
|
payload: KnowledgeFSInitialSourcePreviewPayload | KnowledgeFSInitialWebsiteSourcePreviewPayload,
|
|
is_canceled: Callable[[], bool] | None = None,
|
|
) -> KnowledgeFSInitialSourcePreviewResponse:
|
|
_raise_if_canceled(is_canceled)
|
|
self.require_visible_credential(tenant_id=tenant_id, account=account, payload=payload)
|
|
credentials = DatasourceProviderService().get_datasource_credentials(
|
|
tenant_id=tenant_id,
|
|
provider=payload.provider,
|
|
plugin_id=payload.plugin_id,
|
|
credential_id=payload.credential_id,
|
|
current_user=account,
|
|
)
|
|
if not credentials:
|
|
raise PermissionError("Datasource credential is unavailable")
|
|
provider_type = DatasourceProviderType(payload.kind)
|
|
runtime = DatasourceManager.get_datasource_runtime(
|
|
provider_id=f"{payload.plugin_id}/{payload.provider}",
|
|
datasource_name=payload.datasource,
|
|
tenant_id=tenant_id,
|
|
datasource_type=provider_type,
|
|
)
|
|
runtime.runtime.credentials = credentials
|
|
parameters = dict(payload.parameters)
|
|
if payload.kind == "website_crawl":
|
|
website_runtime = cast(WebsiteCrawlDatasourcePlugin, runtime)
|
|
pages_by_url: dict[str, KnowledgeFSInitialSourcePreviewPageResponse] = {}
|
|
for website_message in website_runtime.get_website_crawl(
|
|
user_id=account.id,
|
|
datasource_parameters=parameters,
|
|
provider_type=website_runtime.datasource_provider_type(),
|
|
):
|
|
_raise_if_canceled(is_canceled)
|
|
for website_page in website_message.result.web_info_list or []:
|
|
pages_by_url[website_page.source_url] = KnowledgeFSInitialSourcePreviewPageResponse(
|
|
description=website_page.description or None,
|
|
source_url=website_page.source_url,
|
|
title=website_page.title or None,
|
|
)
|
|
if len(pages_by_url) >= _MAX_PREVIEW_ITEMS:
|
|
return KnowledgeFSInitialSourcePreviewResponse(
|
|
kind=payload.kind,
|
|
pages=list(pages_by_url.values()),
|
|
)
|
|
return KnowledgeFSInitialSourcePreviewResponse(
|
|
kind=payload.kind,
|
|
pages=list(pages_by_url.values()),
|
|
)
|
|
if payload.kind == "online_document":
|
|
document_runtime = cast(OnlineDocumentDatasourcePlugin, runtime)
|
|
documents: list[KnowledgeFSInitialSourcePreviewDocumentResponse] = []
|
|
for document_message in document_runtime.get_online_document_pages(
|
|
user_id=account.id,
|
|
datasource_parameters=parameters,
|
|
provider_type=document_runtime.datasource_provider_type(),
|
|
):
|
|
for workspace in document_message.result:
|
|
workspace_id = workspace.workspace_id or payload.provider
|
|
for document_page in workspace.pages:
|
|
documents.append(
|
|
KnowledgeFSInitialSourcePreviewDocumentResponse(
|
|
last_edited_time=document_page.last_edited_time,
|
|
name=document_page.page_name,
|
|
page_id=document_page.page_id,
|
|
provider_item_id=json.dumps(
|
|
[workspace_id, document_page.page_id], separators=(",", ":")
|
|
),
|
|
type=document_page.type,
|
|
workspace_id=workspace_id,
|
|
workspace_name=workspace.workspace_name,
|
|
)
|
|
)
|
|
if len(documents) >= _MAX_PREVIEW_ITEMS:
|
|
return KnowledgeFSInitialSourcePreviewResponse(
|
|
documents=documents,
|
|
kind=payload.kind,
|
|
)
|
|
return KnowledgeFSInitialSourcePreviewResponse(documents=documents, kind=payload.kind)
|
|
|
|
drive_runtime = cast(OnlineDriveDatasourcePlugin, runtime)
|
|
files: list[KnowledgeFSInitialSourcePreviewFileResponse] = []
|
|
next_page_parameters = None
|
|
max_keys = parameters.get("max_keys", _MAX_PREVIEW_ITEMS)
|
|
if not isinstance(max_keys, int) or isinstance(max_keys, bool):
|
|
max_keys = _MAX_PREVIEW_ITEMS
|
|
max_keys = min(max(max_keys, 1), _MAX_PREVIEW_ITEMS)
|
|
bucket = parameters.get("bucket")
|
|
prefix = parameters.get("prefix")
|
|
raw_next_page_parameters = parameters.get("next_page_parameters")
|
|
request = OnlineDriveBrowseFilesRequest(
|
|
bucket=bucket if isinstance(bucket, str) else None,
|
|
prefix=prefix if isinstance(prefix, str) else "",
|
|
max_keys=max_keys,
|
|
next_page_parameters=(
|
|
cast(dict[str, Any], raw_next_page_parameters) if isinstance(raw_next_page_parameters, dict) else None
|
|
),
|
|
)
|
|
for drive_message in drive_runtime.online_drive_browse_files(
|
|
user_id=account.id,
|
|
request=request,
|
|
provider_type=drive_runtime.datasource_provider_type(),
|
|
):
|
|
for group in drive_message.result:
|
|
if group.is_truncated and group.next_page_parameters:
|
|
next_page_parameters = group.next_page_parameters
|
|
if group.bucket and not group.files:
|
|
files.append(
|
|
KnowledgeFSInitialSourcePreviewFileResponse(
|
|
bucket=group.bucket,
|
|
id="",
|
|
name=group.bucket,
|
|
provider_item_id=json.dumps([group.bucket, ""], separators=(",", ":")),
|
|
size=0,
|
|
type="bucket",
|
|
)
|
|
)
|
|
for file in group.files:
|
|
files.append(
|
|
KnowledgeFSInitialSourcePreviewFileResponse(
|
|
bucket=group.bucket,
|
|
id=file.id,
|
|
mime_type=file.type if "/" in file.type else None,
|
|
name=file.name,
|
|
provider_item_id=json.dumps([group.bucket or "", file.id], separators=(",", ":")),
|
|
size=file.size,
|
|
type=file.type,
|
|
)
|
|
)
|
|
if len(files) >= _MAX_PREVIEW_ITEMS:
|
|
return KnowledgeFSInitialSourcePreviewResponse(
|
|
files=files,
|
|
kind=payload.kind,
|
|
next_page_parameters=next_page_parameters,
|
|
)
|
|
return KnowledgeFSInitialSourcePreviewResponse(
|
|
files=files,
|
|
kind=payload.kind,
|
|
next_page_parameters=next_page_parameters,
|
|
)
|
|
|
|
|
|
__all__ = ["KnowledgeFSInitialSourcePreviewCanceledError", "KnowledgeFSInitialSourcePreviewService"]
|