dify

dify/api/services/entities/knowledge_entities/knowledge_entities.py

@@ -0,0 +1,169 @@
+from enum import StrEnum
+from typing import Literal
+
+from pydantic import BaseModel
+
+from core.rag.retrieval.retrieval_methods import RetrievalMethod
+
+
+class ParentMode(StrEnum):
+    FULL_DOC = "full-doc"
+    PARAGRAPH = "paragraph"
+
+
+class NotionIcon(BaseModel):
+    type: str
+    url: str | None = None
+    emoji: str | None = None
+
+
+class NotionPage(BaseModel):
+    page_id: str
+    page_name: str
+    page_icon: NotionIcon | None = None
+    type: str
+
+
+class NotionInfo(BaseModel):
+    credential_id: str
+    workspace_id: str
+    pages: list[NotionPage]
+
+
+class WebsiteInfo(BaseModel):
+    provider: str
+    job_id: str
+    urls: list[str]
+    only_main_content: bool = True
+
+
+class FileInfo(BaseModel):
+    file_ids: list[str]
+
+
+class InfoList(BaseModel):
+    data_source_type: Literal["upload_file", "notion_import", "website_crawl"]
+    notion_info_list: list[NotionInfo] | None = None
+    file_info_list: FileInfo | None = None
+    website_info_list: WebsiteInfo | None = None
+
+
+class DataSource(BaseModel):
+    info_list: InfoList
+
+
+class PreProcessingRule(BaseModel):
+    id: str
+    enabled: bool
+
+
+class Segmentation(BaseModel):
+    separator: str = "\n"
+    max_tokens: int
+    chunk_overlap: int = 0
+
+
+class Rule(BaseModel):
+    pre_processing_rules: list[PreProcessingRule] | None = None
+    segmentation: Segmentation | None = None
+    parent_mode: Literal["full-doc", "paragraph"] | None = None
+    subchunk_segmentation: Segmentation | None = None
+
+
+class ProcessRule(BaseModel):
+    mode: Literal["automatic", "custom", "hierarchical"]
+    rules: Rule | None = None
+
+
+class RerankingModel(BaseModel):
+    reranking_provider_name: str | None = None
+    reranking_model_name: str | None = None
+
+
+class WeightVectorSetting(BaseModel):
+    vector_weight: float
+    embedding_provider_name: str
+    embedding_model_name: str
+
+
+class WeightKeywordSetting(BaseModel):
+    keyword_weight: float
+
+
+class WeightModel(BaseModel):
+    weight_type: Literal["semantic_first", "keyword_first", "customized"] | None = None
+    vector_setting: WeightVectorSetting | None = None
+    keyword_setting: WeightKeywordSetting | None = None
+
+
+class RetrievalModel(BaseModel):
+    search_method: RetrievalMethod
+    reranking_enable: bool
+    reranking_model: RerankingModel | None = None
+    reranking_mode: str | None = None
+    top_k: int
+    score_threshold_enabled: bool
+    score_threshold: float | None = None
+    weights: WeightModel | None = None
+
+
+class MetaDataConfig(BaseModel):
+    doc_type: str
+    doc_metadata: dict
+
+
+class KnowledgeConfig(BaseModel):
+    original_document_id: str | None = None
+    duplicate: bool = True
+    indexing_technique: Literal["high_quality", "economy"]
+    data_source: DataSource | None = None
+    process_rule: ProcessRule | None = None
+    retrieval_model: RetrievalModel | None = None
+    doc_form: str = "text_model"
+    doc_language: str = "English"
+    embedding_model: str | None = None
+    embedding_model_provider: str | None = None
+    name: str | None = None
+
+
+class SegmentUpdateArgs(BaseModel):
+    content: str | None = None
+    answer: str | None = None
+    keywords: list[str] | None = None
+    regenerate_child_chunks: bool = False
+    enabled: bool | None = None
+
+
+class ChildChunkUpdateArgs(BaseModel):
+    id: str | None = None
+    content: str
+
+
+class MetadataArgs(BaseModel):
+    type: Literal["string", "number", "time"]
+    name: str
+
+
+class MetadataUpdateArgs(BaseModel):
+    name: str
+    value: str | int | float | None = None
+
+
+class MetadataDetail(BaseModel):
+    id: str
+    name: str
+    value: str | int | float | None = None
+
+
+class DocumentMetadataOperation(BaseModel):
+    document_id: str
+    metadata_list: list[MetadataDetail]
+    partial_update: bool = False
+
+
+class MetadataOperationData(BaseModel):
+    """
+    Metadata operation data
+    """
+
+    operation_data: list[DocumentMetadataOperation]

dify/api/services/entities/knowledge_entities/rag_pipeline_entities.py

@@ -0,0 +1,132 @@
+from typing import Literal
+
+from pydantic import BaseModel, field_validator
+
+from core.rag.retrieval.retrieval_methods import RetrievalMethod
+
+
+class IconInfo(BaseModel):
+    icon: str
+    icon_background: str | None = None
+    icon_type: str | None = None
+    icon_url: str | None = None
+
+
+class PipelineTemplateInfoEntity(BaseModel):
+    name: str
+    description: str
+    icon_info: IconInfo
+
+
+class RagPipelineDatasetCreateEntity(BaseModel):
+    name: str
+    description: str
+    icon_info: IconInfo
+    permission: str
+    partial_member_list: list[str] | None = None
+    yaml_content: str | None = None
+
+
+class RerankingModelConfig(BaseModel):
+    """
+    Reranking Model Config.
+    """
+
+    reranking_provider_name: str | None = ""
+    reranking_model_name: str | None = ""
+
+
+class VectorSetting(BaseModel):
+    """
+    Vector Setting.
+    """
+
+    vector_weight: float
+    embedding_provider_name: str
+    embedding_model_name: str
+
+
+class KeywordSetting(BaseModel):
+    """
+    Keyword Setting.
+    """
+
+    keyword_weight: float
+
+
+class WeightedScoreConfig(BaseModel):
+    """
+    Weighted score Config.
+    """
+
+    vector_setting: VectorSetting | None
+    keyword_setting: KeywordSetting | None
+
+
+class EmbeddingSetting(BaseModel):
+    """
+    Embedding Setting.
+    """
+
+    embedding_provider_name: str
+    embedding_model_name: str
+
+
+class EconomySetting(BaseModel):
+    """
+    Economy Setting.
+    """
+
+    keyword_number: int
+
+
+class RetrievalSetting(BaseModel):
+    """
+    Retrieval Setting.
+    """
+
+    search_method: RetrievalMethod
+    top_k: int
+    score_threshold: float | None = 0.5
+    score_threshold_enabled: bool = False
+    reranking_mode: str | None = "reranking_model"
+    reranking_enable: bool | None = True
+    reranking_model: RerankingModelConfig | None = None
+    weights: WeightedScoreConfig | None = None
+
+
+class IndexMethod(BaseModel):
+    """
+    Knowledge Index Setting.
+    """
+
+    indexing_technique: Literal["high_quality", "economy"]
+    embedding_setting: EmbeddingSetting
+    economy_setting: EconomySetting
+
+
+class KnowledgeConfiguration(BaseModel):
+    """
+    Knowledge Base Configuration.
+    """
+
+    chunk_structure: str
+    indexing_technique: Literal["high_quality", "economy"]
+    embedding_model_provider: str = ""
+    embedding_model: str = ""
+    keyword_number: int | None = 10
+    retrieval_model: RetrievalSetting
+
+    @field_validator("embedding_model_provider", mode="before")
+    @classmethod
+    def validate_embedding_model_provider(cls, v):
+        if v is None:
+            return ""
+        return v
+
+    @field_validator("embedding_model", mode="before")
+    @classmethod
+    def validate_embedding_model(cls, v):
+        if v is None:
+            return ""
+        return v