dify/api/services/entities/knowledge_entities/rag_pipeline_entities.py
2025-06-03 19:02:57 +08:00

117 lines
2.3 KiB
Python

from typing import Literal, Optional
from pydantic import BaseModel
class IconInfo(BaseModel):
icon: str
icon_background: Optional[str] = None
icon_type: Optional[str] = None
icon_url: Optional[str] = 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: Optional[list[str]] = None
yaml_content: Optional[str] = None
class RerankingModelConfig(BaseModel):
"""
Reranking Model Config.
"""
reranking_provider_name: str
reranking_model_name: str
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
keyword_setting: KeywordSetting
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: Literal["semantic_search", "fulltext_search", "keyword_search", "hybrid_search"]
top_k: int
score_threshold: Optional[float] = 0.5
score_threshold_enabled: bool = False
reranking_mode: str = "reranking_model"
reranking_enable: bool = True
reranking_model: Optional[RerankingModelConfig] = None
weights: Optional[WeightedScoreConfig] = 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: Optional[str] = ""
embedding_model: Optional[str] = ""
keyword_number: Optional[int] = 10
retrieval_model: RetrievalSetting