mirror of
https://github.com/langgenius/dify.git
synced 2026-09-28 14:23:33 +08:00
476 lines
16 KiB
Python
476 lines
16 KiB
Python
from collections.abc import Sequence
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
import sqlalchemy as sa
|
|
from pydantic import Field, field_validator
|
|
from sqlalchemy import func, select
|
|
from sqlalchemy.orm import Session
|
|
|
|
from fields.base import ResponseModel
|
|
from libs.helper import to_timestamp
|
|
from models.account import Account
|
|
from models.dataset import (
|
|
AppDatasetJoin,
|
|
Dataset,
|
|
DatasetMetadata,
|
|
Document,
|
|
ExternalKnowledgeApis,
|
|
ExternalKnowledgeBindings,
|
|
Pipeline,
|
|
)
|
|
from models.model import App, Tag, TagBinding
|
|
|
|
|
|
class DatasetMetadataResponse(ResponseModel):
|
|
id: str
|
|
type: str
|
|
name: str
|
|
|
|
|
|
class DatasetMetadataListItemResponse(ResponseModel):
|
|
id: str
|
|
name: str
|
|
type: str
|
|
count: int = 0
|
|
|
|
|
|
class DatasetMetadataListResponse(ResponseModel):
|
|
doc_metadata: list[DatasetMetadataListItemResponse]
|
|
built_in_field_enabled: bool
|
|
|
|
|
|
class DatasetMetadataBuiltInFieldResponse(ResponseModel):
|
|
name: str
|
|
type: str
|
|
|
|
|
|
class DatasetMetadataBuiltInFieldsResponse(ResponseModel):
|
|
fields: list[DatasetMetadataBuiltInFieldResponse]
|
|
|
|
|
|
class DatasetMetadataActionResponse(ResponseModel):
|
|
result: str = Field(description="Operation result.")
|
|
|
|
|
|
class DatasetRerankingModelResponse(ResponseModel):
|
|
reranking_provider_name: str | None = None
|
|
reranking_model_name: str | None = None
|
|
|
|
|
|
class DatasetKeywordSettingResponse(ResponseModel):
|
|
keyword_weight: float | None = None
|
|
|
|
|
|
class DatasetVectorSettingResponse(ResponseModel):
|
|
vector_weight: float | None = None
|
|
embedding_model_name: str | None = None
|
|
embedding_provider_name: str | None = None
|
|
|
|
|
|
class DatasetWeightedScoreResponse(ResponseModel):
|
|
weight_type: str | None = None
|
|
keyword_setting: DatasetKeywordSettingResponse = Field(
|
|
default_factory=DatasetKeywordSettingResponse,
|
|
description="Keyword search weight settings.",
|
|
)
|
|
vector_setting: DatasetVectorSettingResponse = Field(
|
|
default_factory=DatasetVectorSettingResponse,
|
|
description="Semantic search weight settings.",
|
|
)
|
|
|
|
@field_validator("keyword_setting", "vector_setting", mode="before")
|
|
@classmethod
|
|
def _expand_null_nested(cls, value: object) -> object:
|
|
return {} if value is None else value
|
|
|
|
|
|
class DatasetRetrievalModelResponse(ResponseModel):
|
|
search_method: str
|
|
reranking_enable: bool
|
|
reranking_mode: str | None = None
|
|
reranking_model: DatasetRerankingModelResponse = Field(
|
|
default_factory=DatasetRerankingModelResponse,
|
|
description="Reranking model configuration.",
|
|
)
|
|
weights: DatasetWeightedScoreResponse | None = None
|
|
top_k: int
|
|
score_threshold_enabled: bool
|
|
score_threshold: float | None = None
|
|
|
|
@field_validator("reranking_model", mode="before")
|
|
@classmethod
|
|
def _expand_null_nested(cls, value: object) -> object:
|
|
return {} if value is None else value
|
|
|
|
|
|
class DatasetSummaryIndexSettingResponse(ResponseModel):
|
|
enable: bool | None = None
|
|
model_name: str | None = None
|
|
model_provider_name: str | None = None
|
|
summary_prompt: str | None = None
|
|
|
|
|
|
class DatasetTagResponse(ResponseModel):
|
|
id: str
|
|
name: str
|
|
type: str
|
|
|
|
|
|
class DatasetExternalKnowledgeInfoResponse(ResponseModel):
|
|
external_knowledge_id: str | None = None
|
|
external_knowledge_api_id: str | None = None
|
|
external_knowledge_api_name: str | None = None
|
|
external_knowledge_api_endpoint: str | None = None
|
|
|
|
|
|
class DatasetExternalRetrievalModelResponse(ResponseModel):
|
|
top_k: int
|
|
score_threshold: float | None = None
|
|
score_threshold_enabled: bool | None = None
|
|
|
|
|
|
class DatasetDocMetadataResponse(ResponseModel):
|
|
id: str
|
|
name: str
|
|
type: str
|
|
|
|
|
|
class DatasetIconInfoResponse(ResponseModel):
|
|
icon_type: str | None = None
|
|
icon: str | None = None
|
|
icon_background: str | None = None
|
|
icon_url: str | None = None
|
|
|
|
|
|
class DatasetDetailResponse(ResponseModel):
|
|
id: str
|
|
name: str
|
|
description: str | None
|
|
provider: str
|
|
permission: str
|
|
data_source_type: str | None
|
|
indexing_technique: str | None
|
|
app_count: int
|
|
document_count: int
|
|
word_count: int
|
|
created_by: str
|
|
author_name: str | None
|
|
created_at: int
|
|
updated_by: str | None
|
|
updated_at: int
|
|
embedding_model: str | None
|
|
embedding_model_provider: str | None
|
|
embedding_available: bool | None = None
|
|
retrieval_model_dict: DatasetRetrievalModelResponse = Field(
|
|
description="Retrieval configuration for the knowledge base."
|
|
)
|
|
summary_index_setting: DatasetSummaryIndexSettingResponse = Field(
|
|
default_factory=DatasetSummaryIndexSettingResponse,
|
|
description="Summary index configuration.",
|
|
)
|
|
tags: list[DatasetTagResponse]
|
|
doc_form: str | None
|
|
external_knowledge_info: DatasetExternalKnowledgeInfoResponse = Field(
|
|
default_factory=DatasetExternalKnowledgeInfoResponse,
|
|
description=(
|
|
"Connection details for external knowledge bases. Populated when `provider` is `external`; otherwise "
|
|
"its properties are `null`."
|
|
),
|
|
)
|
|
external_retrieval_model: DatasetExternalRetrievalModelResponse | None
|
|
doc_metadata: list[DatasetDocMetadataResponse]
|
|
built_in_field_enabled: bool
|
|
pipeline_id: str | None
|
|
runtime_mode: str | None
|
|
chunk_structure: str | None
|
|
icon_info: DatasetIconInfoResponse = Field(
|
|
default_factory=DatasetIconInfoResponse,
|
|
description="Icon display configuration for the knowledge base.",
|
|
)
|
|
is_published: bool
|
|
total_documents: int
|
|
total_available_documents: int
|
|
enable_api: bool
|
|
is_multimodal: bool
|
|
permission_keys: list[str] = Field(default_factory=list)
|
|
maintainer: str | None = None
|
|
|
|
@field_validator("created_at", "updated_at", mode="before")
|
|
@classmethod
|
|
def _normalize_timestamp(cls, value: datetime | int | None) -> int | None:
|
|
return to_timestamp(value)
|
|
|
|
@field_validator("summary_index_setting", "external_knowledge_info", "icon_info", mode="before")
|
|
@classmethod
|
|
def _expand_null_nested(cls, value: object) -> object:
|
|
return {} if value is None else value
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DatasetDetailPrefetch:
|
|
"""Batch-loaded values backing the session-scoped fields of a page of datasets.
|
|
|
|
Building a detail response for a single dataset issues about ten queries for
|
|
that row alone, so a list endpoint costs one query per dataset per field.
|
|
Loading the same values for every dataset on the page up front keeps the
|
|
query count bounded regardless of the page size.
|
|
"""
|
|
|
|
app_counts: dict[str, int] = field(default_factory=dict)
|
|
document_counts: dict[str, int] = field(default_factory=dict)
|
|
word_counts: dict[str, int] = field(default_factory=dict)
|
|
available_document_counts: dict[str, int] = field(default_factory=dict)
|
|
author_names: dict[str, str] = field(default_factory=dict)
|
|
tags: dict[tuple[str, str], list[Tag]] = field(default_factory=dict)
|
|
doc_forms: dict[tuple[str, str], str | None] = field(default_factory=dict)
|
|
external_knowledge_infos: dict[tuple[str, str], dict[str, Any]] = field(default_factory=dict)
|
|
doc_metadata: dict[str, list[DatasetMetadata]] = field(default_factory=dict)
|
|
pipeline_published: dict[str, bool] = field(default_factory=dict)
|
|
|
|
|
|
def build_dataset_detail_prefetch(datasets: Sequence[Dataset], *, session: Session) -> DatasetDetailPrefetch:
|
|
"""Load every session-scoped detail field for the given datasets in a fixed number of queries."""
|
|
dataset_ids = [dataset.id for dataset in datasets]
|
|
if not dataset_ids:
|
|
return DatasetDetailPrefetch()
|
|
|
|
tenant_ids = {dataset.tenant_id for dataset in datasets}
|
|
|
|
app_counts = dict(
|
|
session.execute(
|
|
select(AppDatasetJoin.dataset_id, func.count(AppDatasetJoin.id))
|
|
.where(AppDatasetJoin.dataset_id.in_(dataset_ids), App.id == AppDatasetJoin.app_id)
|
|
.group_by(AppDatasetJoin.dataset_id)
|
|
)
|
|
.tuples()
|
|
.all()
|
|
)
|
|
|
|
document_counts: dict[str, int] = {}
|
|
word_counts: dict[str, int] = {}
|
|
available_document_counts: dict[str, int] = {}
|
|
available_document_case = sa.case(
|
|
(
|
|
sa.and_(
|
|
Document.indexing_status == "completed",
|
|
Document.enabled == True,
|
|
Document.archived == False,
|
|
),
|
|
1,
|
|
),
|
|
else_=0,
|
|
)
|
|
document_rows = session.execute(
|
|
select(
|
|
Document.dataset_id,
|
|
func.count(Document.id),
|
|
func.coalesce(func.sum(Document.word_count), 0),
|
|
func.coalesce(func.sum(available_document_case), 0),
|
|
)
|
|
.where(Document.dataset_id.in_(dataset_ids))
|
|
.group_by(Document.dataset_id)
|
|
).all()
|
|
for dataset_id, document_count, word_count, available_count in document_rows:
|
|
document_counts[dataset_id] = document_count
|
|
word_counts[dataset_id] = word_count
|
|
available_document_counts[dataset_id] = available_count
|
|
|
|
# doc_form only falls back to the documents table for datasets without a
|
|
# chunk structure of their own.
|
|
doc_forms: dict[tuple[str, str], str | None] = {}
|
|
doc_form_dataset_ids = [dataset.id for dataset in datasets if not dataset.chunk_structure]
|
|
if doc_form_dataset_ids:
|
|
doc_forms = {
|
|
(dataset_id, tenant_id): doc_form
|
|
for dataset_id, tenant_id, doc_form in session.execute(
|
|
select(Document.dataset_id, Document.tenant_id, func.min(Document.doc_form))
|
|
.where(Document.dataset_id.in_(doc_form_dataset_ids), Document.tenant_id.in_(tenant_ids))
|
|
.group_by(Document.dataset_id, Document.tenant_id)
|
|
).all()
|
|
}
|
|
|
|
author_ids = {dataset.created_by for dataset in datasets if dataset.created_by}
|
|
author_names = dict(
|
|
session.execute(select(Account.id, Account.name).where(Account.id.in_(author_ids))).tuples().all()
|
|
)
|
|
|
|
tags: dict[tuple[str, str], list[Tag]] = {}
|
|
tag_rows = session.execute(
|
|
select(TagBinding.target_id, TagBinding.tenant_id, Tag)
|
|
.join(TagBinding, Tag.id == TagBinding.tag_id)
|
|
.where(
|
|
TagBinding.target_id.in_(dataset_ids),
|
|
TagBinding.tenant_id.in_(tenant_ids),
|
|
Tag.tenant_id == TagBinding.tenant_id,
|
|
Tag.type == "knowledge",
|
|
)
|
|
).all()
|
|
for target_id, tenant_id, tag in tag_rows:
|
|
tags.setdefault((target_id, tenant_id), []).append(tag)
|
|
|
|
doc_metadata: dict[str, list[DatasetMetadata]] = {}
|
|
for dataset_metadata in session.scalars(
|
|
select(DatasetMetadata).where(DatasetMetadata.dataset_id.in_(dataset_ids))
|
|
).all():
|
|
doc_metadata.setdefault(dataset_metadata.dataset_id, []).append(dataset_metadata)
|
|
|
|
pipeline_ids = {dataset.pipeline_id for dataset in datasets if dataset.pipeline_id}
|
|
pipeline_published: dict[str, bool] = {}
|
|
if pipeline_ids:
|
|
pipeline_published = dict(
|
|
session.execute(select(Pipeline.id, Pipeline.is_published).where(Pipeline.id.in_(pipeline_ids)))
|
|
.tuples()
|
|
.all()
|
|
)
|
|
|
|
return DatasetDetailPrefetch(
|
|
app_counts=app_counts,
|
|
document_counts=document_counts,
|
|
word_counts=word_counts,
|
|
available_document_counts=available_document_counts,
|
|
author_names=author_names,
|
|
tags=tags,
|
|
doc_forms=doc_forms,
|
|
external_knowledge_infos=_load_external_knowledge_infos(datasets, session=session),
|
|
doc_metadata=doc_metadata,
|
|
pipeline_published=pipeline_published,
|
|
)
|
|
|
|
|
|
def _load_external_knowledge_infos(
|
|
datasets: Sequence[Dataset], *, session: Session
|
|
) -> dict[tuple[str, str], dict[str, Any]]:
|
|
external_datasets = [dataset for dataset in datasets if dataset.provider == "external"]
|
|
if not external_datasets:
|
|
return {}
|
|
|
|
external_dataset_ids = [dataset.id for dataset in external_datasets]
|
|
external_tenant_ids = {dataset.tenant_id for dataset in external_datasets}
|
|
bindings = session.scalars(
|
|
select(ExternalKnowledgeBindings).where(
|
|
ExternalKnowledgeBindings.dataset_id.in_(external_dataset_ids),
|
|
ExternalKnowledgeBindings.tenant_id.in_(external_tenant_ids),
|
|
)
|
|
).all()
|
|
if not bindings:
|
|
return {}
|
|
|
|
apis = {
|
|
(api.id, api.tenant_id): api
|
|
for api in session.scalars(
|
|
select(ExternalKnowledgeApis).where(
|
|
ExternalKnowledgeApis.id.in_({binding.external_knowledge_api_id for binding in bindings}),
|
|
ExternalKnowledgeApis.tenant_id.in_(external_tenant_ids),
|
|
)
|
|
).all()
|
|
}
|
|
|
|
infos: dict[tuple[str, str], dict[str, Any]] = {}
|
|
for binding in bindings:
|
|
info = Dataset.build_external_knowledge_info(
|
|
binding, apis.get((binding.external_knowledge_api_id, binding.tenant_id))
|
|
)
|
|
if info is not None:
|
|
infos[(binding.dataset_id, binding.tenant_id)] = info
|
|
return infos
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DatasetDetailResponseSource:
|
|
"""Expose session-backed dataset fields during response validation.
|
|
|
|
With a prefetch the values come from the batch load for the whole page;
|
|
without one each field is resolved against the session on access.
|
|
"""
|
|
|
|
dataset: Any
|
|
session: Session
|
|
prefetch: DatasetDetailPrefetch | None = None
|
|
|
|
@property
|
|
def _prefetch_key(self) -> tuple[str, str]:
|
|
return (self.dataset.id, self.dataset.tenant_id)
|
|
|
|
@property
|
|
def app_count(self) -> int:
|
|
if self.prefetch is not None:
|
|
return self.prefetch.app_counts.get(self.dataset.id, 0)
|
|
return self.dataset.get_app_count(session=self.session)
|
|
|
|
@property
|
|
def document_count(self) -> int:
|
|
if self.prefetch is not None:
|
|
return self.prefetch.document_counts.get(self.dataset.id, 0)
|
|
return self.dataset.get_document_count(session=self.session)
|
|
|
|
@property
|
|
def word_count(self) -> int:
|
|
if self.prefetch is not None:
|
|
return self.prefetch.word_counts.get(self.dataset.id, 0)
|
|
return self.dataset.get_word_count(session=self.session)
|
|
|
|
@property
|
|
def author_name(self) -> str | None:
|
|
if self.prefetch is not None:
|
|
return self.prefetch.author_names.get(self.dataset.created_by)
|
|
return self.dataset.get_author_name(session=self.session)
|
|
|
|
@property
|
|
def tags(self) -> Any:
|
|
if self.prefetch is not None:
|
|
return self.prefetch.tags.get(self._prefetch_key, [])
|
|
return self.dataset.get_tags(session=self.session)
|
|
|
|
@property
|
|
def doc_form(self) -> str | None:
|
|
if self.prefetch is not None:
|
|
if self.dataset.chunk_structure:
|
|
return self.dataset.chunk_structure
|
|
return self.prefetch.doc_forms.get(self._prefetch_key)
|
|
return self.dataset.get_doc_form(session=self.session)
|
|
|
|
@property
|
|
def external_knowledge_info(self) -> Any:
|
|
if self.prefetch is not None:
|
|
if self.dataset.provider != "external":
|
|
return None
|
|
return self.prefetch.external_knowledge_infos.get(self._prefetch_key)
|
|
return self.dataset.get_external_knowledge_info(session=self.session)
|
|
|
|
@property
|
|
def doc_metadata(self) -> Any:
|
|
if self.prefetch is not None:
|
|
return self.dataset.build_doc_metadata(self.prefetch.doc_metadata.get(self.dataset.id, []))
|
|
return self.dataset.get_doc_metadata(session=self.session)
|
|
|
|
@property
|
|
def is_published(self) -> bool:
|
|
if self.prefetch is not None:
|
|
if not self.dataset.pipeline_id:
|
|
return False
|
|
return self.prefetch.pipeline_published.get(self.dataset.pipeline_id, False)
|
|
return self.dataset.get_is_published(session=self.session)
|
|
|
|
@property
|
|
def total_documents(self) -> int:
|
|
if self.prefetch is not None:
|
|
return self.prefetch.document_counts.get(self.dataset.id, 0)
|
|
return self.dataset.get_total_documents(session=self.session)
|
|
|
|
@property
|
|
def total_available_documents(self) -> int:
|
|
if self.prefetch is not None:
|
|
return self.prefetch.available_document_counts.get(self.dataset.id, 0)
|
|
return self.dataset.get_total_available_documents(session=self.session)
|
|
|
|
def __getattr__(self, name: str) -> Any:
|
|
return getattr(self.dataset, name) # guard-ignore: no-new-getattr -- delegates model fields
|
|
|
|
|
|
def dataset_detail_response_source(
|
|
dataset: Any, *, session: Session, prefetch: DatasetDetailPrefetch | None = None
|
|
) -> DatasetDetailResponseSource:
|
|
return DatasetDetailResponseSource(dataset=dataset, session=session, prefetch=prefetch)
|