diff --git a/.gitattributes b/.gitattributes index 795574b32ce3a0b061b5bf1d6eebe5be2f50b20b..80ff363ec6a36b16cfeb9d2f9b8a1fa0f8227a1a 100644 --- a/.gitattributes +++ b/.gitattributes @@ -6,3 +6,9 @@ *.sh text eol=lf api/tests/integration_tests/model_runtime/assets/audio.mp3 filter=lfs diff=lfs merge=lfs -text +api/core/tools/docs/images/index/image-1.png filter=lfs diff=lfs merge=lfs -text +api/core/tools/docs/images/index/image-2.png filter=lfs diff=lfs merge=lfs -text +api/core/tools/docs/images/index/image.png filter=lfs diff=lfs merge=lfs -text +api/core/tools/provider/builtin/comfyui/_assets/icon.png filter=lfs diff=lfs merge=lfs -text +api/core/tools/provider/builtin/dalle/_assets/icon.png filter=lfs diff=lfs merge=lfs -text +api/core/tools/provider/builtin/wecom/_assets/icon.png filter=lfs diff=lfs merge=lfs -text diff --git a/api/core/__init__.py b/api/core/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6eaea7b1c8419f6875babcd591f25088a01a5527 --- /dev/null +++ b/api/core/__init__.py @@ -0,0 +1 @@ +import core.moderation.base diff --git a/api/core/hosting_configuration.py b/api/core/hosting_configuration.py new file mode 100644 index 0000000000000000000000000000000000000000..f9fb7275f3624ffc88e8d6b82d7f1c52f5748a69 --- /dev/null +++ b/api/core/hosting_configuration.py @@ -0,0 +1,255 @@ +from typing import Optional + +from flask import Flask +from pydantic import BaseModel + +from configs import dify_config +from core.entities.provider_entities import QuotaUnit, RestrictModel +from core.model_runtime.entities.model_entities import ModelType +from models.provider import ProviderQuotaType + + +class HostingQuota(BaseModel): + quota_type: ProviderQuotaType + restrict_models: list[RestrictModel] = [] + + +class TrialHostingQuota(HostingQuota): + quota_type: ProviderQuotaType = ProviderQuotaType.TRIAL + quota_limit: int = 0 + """Quota limit for the hosting provider models. -1 means unlimited.""" + + +class PaidHostingQuota(HostingQuota): + quota_type: ProviderQuotaType = ProviderQuotaType.PAID + + +class FreeHostingQuota(HostingQuota): + quota_type: ProviderQuotaType = ProviderQuotaType.FREE + + +class HostingProvider(BaseModel): + enabled: bool = False + credentials: Optional[dict] = None + quota_unit: Optional[QuotaUnit] = None + quotas: list[HostingQuota] = [] + + +class HostedModerationConfig(BaseModel): + enabled: bool = False + providers: list[str] = [] + + +class HostingConfiguration: + provider_map: dict[str, HostingProvider] = {} + moderation_config: Optional[HostedModerationConfig] = None + + def init_app(self, app: Flask) -> None: + if dify_config.EDITION != "CLOUD": + return + + self.provider_map["azure_openai"] = self.init_azure_openai() + self.provider_map["openai"] = self.init_openai() + self.provider_map["anthropic"] = self.init_anthropic() + self.provider_map["minimax"] = self.init_minimax() + self.provider_map["spark"] = self.init_spark() + self.provider_map["zhipuai"] = self.init_zhipuai() + + self.moderation_config = self.init_moderation_config() + + @staticmethod + def init_azure_openai() -> HostingProvider: + quota_unit = QuotaUnit.TIMES + if dify_config.HOSTED_AZURE_OPENAI_ENABLED: + credentials = { + "openai_api_key": dify_config.HOSTED_AZURE_OPENAI_API_KEY, + "openai_api_base": dify_config.HOSTED_AZURE_OPENAI_API_BASE, + "base_model_name": "gpt-35-turbo", + } + + quotas: list[HostingQuota] = [] + hosted_quota_limit = dify_config.HOSTED_AZURE_OPENAI_QUOTA_LIMIT + trial_quota = TrialHostingQuota( + quota_limit=hosted_quota_limit, + restrict_models=[ + RestrictModel(model="gpt-4", base_model_name="gpt-4", model_type=ModelType.LLM), + RestrictModel(model="gpt-4o", base_model_name="gpt-4o", model_type=ModelType.LLM), + RestrictModel(model="gpt-4o-mini", base_model_name="gpt-4o-mini", model_type=ModelType.LLM), + RestrictModel(model="gpt-4-32k", base_model_name="gpt-4-32k", model_type=ModelType.LLM), + RestrictModel( + model="gpt-4-1106-preview", base_model_name="gpt-4-1106-preview", model_type=ModelType.LLM + ), + RestrictModel( + model="gpt-4-vision-preview", base_model_name="gpt-4-vision-preview", model_type=ModelType.LLM + ), + RestrictModel(model="gpt-35-turbo", base_model_name="gpt-35-turbo", model_type=ModelType.LLM), + RestrictModel( + model="gpt-35-turbo-1106", base_model_name="gpt-35-turbo-1106", model_type=ModelType.LLM + ), + RestrictModel( + model="gpt-35-turbo-instruct", base_model_name="gpt-35-turbo-instruct", model_type=ModelType.LLM + ), + RestrictModel( + model="gpt-35-turbo-16k", base_model_name="gpt-35-turbo-16k", model_type=ModelType.LLM + ), + RestrictModel( + model="text-davinci-003", base_model_name="text-davinci-003", model_type=ModelType.LLM + ), + RestrictModel( + model="text-embedding-ada-002", + base_model_name="text-embedding-ada-002", + model_type=ModelType.TEXT_EMBEDDING, + ), + RestrictModel( + model="text-embedding-3-small", + base_model_name="text-embedding-3-small", + model_type=ModelType.TEXT_EMBEDDING, + ), + RestrictModel( + model="text-embedding-3-large", + base_model_name="text-embedding-3-large", + model_type=ModelType.TEXT_EMBEDDING, + ), + ], + ) + quotas.append(trial_quota) + + return HostingProvider(enabled=True, credentials=credentials, quota_unit=quota_unit, quotas=quotas) + + return HostingProvider( + enabled=False, + quota_unit=quota_unit, + ) + + def init_openai(self) -> HostingProvider: + quota_unit = QuotaUnit.CREDITS + quotas: list[HostingQuota] = [] + + if dify_config.HOSTED_OPENAI_TRIAL_ENABLED: + hosted_quota_limit = dify_config.HOSTED_OPENAI_QUOTA_LIMIT + trial_models = self.parse_restrict_models_from_env("HOSTED_OPENAI_TRIAL_MODELS") + trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit, restrict_models=trial_models) + quotas.append(trial_quota) + + if dify_config.HOSTED_OPENAI_PAID_ENABLED: + paid_models = self.parse_restrict_models_from_env("HOSTED_OPENAI_PAID_MODELS") + paid_quota = PaidHostingQuota(restrict_models=paid_models) + quotas.append(paid_quota) + + if len(quotas) > 0: + credentials = { + "openai_api_key": dify_config.HOSTED_OPENAI_API_KEY, + } + + if dify_config.HOSTED_OPENAI_API_BASE: + credentials["openai_api_base"] = dify_config.HOSTED_OPENAI_API_BASE + + if dify_config.HOSTED_OPENAI_API_ORGANIZATION: + credentials["openai_organization"] = dify_config.HOSTED_OPENAI_API_ORGANIZATION + + return HostingProvider(enabled=True, credentials=credentials, quota_unit=quota_unit, quotas=quotas) + + return HostingProvider( + enabled=False, + quota_unit=quota_unit, + ) + + @staticmethod + def init_anthropic() -> HostingProvider: + quota_unit = QuotaUnit.TOKENS + quotas: list[HostingQuota] = [] + + if dify_config.HOSTED_ANTHROPIC_TRIAL_ENABLED: + hosted_quota_limit = dify_config.HOSTED_ANTHROPIC_QUOTA_LIMIT + trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit) + quotas.append(trial_quota) + + if dify_config.HOSTED_ANTHROPIC_PAID_ENABLED: + paid_quota = PaidHostingQuota() + quotas.append(paid_quota) + + if len(quotas) > 0: + credentials = { + "anthropic_api_key": dify_config.HOSTED_ANTHROPIC_API_KEY, + } + + if dify_config.HOSTED_ANTHROPIC_API_BASE: + credentials["anthropic_api_url"] = dify_config.HOSTED_ANTHROPIC_API_BASE + + return HostingProvider(enabled=True, credentials=credentials, quota_unit=quota_unit, quotas=quotas) + + return HostingProvider( + enabled=False, + quota_unit=quota_unit, + ) + + @staticmethod + def init_minimax() -> HostingProvider: + quota_unit = QuotaUnit.TOKENS + if dify_config.HOSTED_MINIMAX_ENABLED: + quotas: list[HostingQuota] = [FreeHostingQuota()] + + return HostingProvider( + enabled=True, + credentials=None, # use credentials from the provider + quota_unit=quota_unit, + quotas=quotas, + ) + + return HostingProvider( + enabled=False, + quota_unit=quota_unit, + ) + + @staticmethod + def init_spark() -> HostingProvider: + quota_unit = QuotaUnit.TOKENS + if dify_config.HOSTED_SPARK_ENABLED: + quotas: list[HostingQuota] = [FreeHostingQuota()] + + return HostingProvider( + enabled=True, + credentials=None, # use credentials from the provider + quota_unit=quota_unit, + quotas=quotas, + ) + + return HostingProvider( + enabled=False, + quota_unit=quota_unit, + ) + + @staticmethod + def init_zhipuai() -> HostingProvider: + quota_unit = QuotaUnit.TOKENS + if dify_config.HOSTED_ZHIPUAI_ENABLED: + quotas: list[HostingQuota] = [FreeHostingQuota()] + + return HostingProvider( + enabled=True, + credentials=None, # use credentials from the provider + quota_unit=quota_unit, + quotas=quotas, + ) + + return HostingProvider( + enabled=False, + quota_unit=quota_unit, + ) + + @staticmethod + def init_moderation_config() -> HostedModerationConfig: + if dify_config.HOSTED_MODERATION_ENABLED and dify_config.HOSTED_MODERATION_PROVIDERS: + return HostedModerationConfig(enabled=True, providers=dify_config.HOSTED_MODERATION_PROVIDERS.split(",")) + + return HostedModerationConfig(enabled=False) + + @staticmethod + def parse_restrict_models_from_env(env_var: str) -> list[RestrictModel]: + models_str = dify_config.model_dump().get(env_var) + models_list = models_str.split(",") if models_str else [] + return [ + RestrictModel(model=model_name.strip(), model_type=ModelType.LLM) + for model_name in models_list + if model_name.strip() + ] diff --git a/api/core/indexing_runner.py b/api/core/indexing_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..1bc4baf9c0e9e760a8162de293588bdc1ff88a37 --- /dev/null +++ b/api/core/indexing_runner.py @@ -0,0 +1,754 @@ +import concurrent.futures +import datetime +import json +import logging +import re +import threading +import time +import uuid +from typing import Any, Optional, cast + +from flask import current_app +from flask_login import current_user # type: ignore +from sqlalchemy.orm.exc import ObjectDeletedError + +from configs import dify_config +from core.entities.knowledge_entities import IndexingEstimate, PreviewDetail, QAPreviewDetail +from core.errors.error import ProviderTokenNotInitError +from core.model_manager import ModelInstance, ModelManager +from core.model_runtime.entities.model_entities import ModelType +from core.rag.cleaner.clean_processor import CleanProcessor +from core.rag.datasource.keyword.keyword_factory import Keyword +from core.rag.docstore.dataset_docstore import DatasetDocumentStore +from core.rag.extractor.entity.extract_setting import ExtractSetting +from core.rag.index_processor.constant.index_type import IndexType +from core.rag.index_processor.index_processor_base import BaseIndexProcessor +from core.rag.index_processor.index_processor_factory import IndexProcessorFactory +from core.rag.models.document import ChildDocument, Document +from core.rag.splitter.fixed_text_splitter import ( + EnhanceRecursiveCharacterTextSplitter, + FixedRecursiveCharacterTextSplitter, +) +from core.rag.splitter.text_splitter import TextSplitter +from core.tools.utils.web_reader_tool import get_image_upload_file_ids +from extensions.ext_database import db +from extensions.ext_redis import redis_client +from extensions.ext_storage import storage +from libs import helper +from models.dataset import ChildChunk, Dataset, DatasetProcessRule, DocumentSegment +from models.dataset import Document as DatasetDocument +from models.model import UploadFile +from services.feature_service import FeatureService + + +class IndexingRunner: + def __init__(self): + self.storage = storage + self.model_manager = ModelManager() + + def run(self, dataset_documents: list[DatasetDocument]): + """Run the indexing process.""" + for dataset_document in dataset_documents: + try: + # get dataset + dataset = Dataset.query.filter_by(id=dataset_document.dataset_id).first() + + if not dataset: + raise ValueError("no dataset found") + + # get the process rule + processing_rule = ( + db.session.query(DatasetProcessRule) + .filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id) + .first() + ) + if not processing_rule: + raise ValueError("no process rule found") + index_type = dataset_document.doc_form + index_processor = IndexProcessorFactory(index_type).init_index_processor() + # extract + text_docs = self._extract(index_processor, dataset_document, processing_rule.to_dict()) + + # transform + documents = self._transform( + index_processor, dataset, text_docs, dataset_document.doc_language, processing_rule.to_dict() + ) + # save segment + self._load_segments(dataset, dataset_document, documents) + + # load + self._load( + index_processor=index_processor, + dataset=dataset, + dataset_document=dataset_document, + documents=documents, + ) + except DocumentIsPausedError: + raise DocumentIsPausedError("Document paused, document id: {}".format(dataset_document.id)) + except ProviderTokenNotInitError as e: + dataset_document.indexing_status = "error" + dataset_document.error = str(e.description) + dataset_document.stopped_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None) + db.session.commit() + except ObjectDeletedError: + logging.warning("Document deleted, document id: {}".format(dataset_document.id)) + except Exception as e: + logging.exception("consume document failed") + dataset_document.indexing_status = "error" + dataset_document.error = str(e) + dataset_document.stopped_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None) + db.session.commit() + + def run_in_splitting_status(self, dataset_document: DatasetDocument): + """Run the indexing process when the index_status is splitting.""" + try: + # get dataset + dataset = Dataset.query.filter_by(id=dataset_document.dataset_id).first() + + if not dataset: + raise ValueError("no dataset found") + + # get exist document_segment list and delete + document_segments = DocumentSegment.query.filter_by( + dataset_id=dataset.id, document_id=dataset_document.id + ).all() + + for document_segment in document_segments: + db.session.delete(document_segment) + if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX: + # delete child chunks + db.session.query(ChildChunk).filter(ChildChunk.segment_id == document_segment.id).delete() + db.session.commit() + # get the process rule + processing_rule = ( + db.session.query(DatasetProcessRule) + .filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id) + .first() + ) + if not processing_rule: + raise ValueError("no process rule found") + + index_type = dataset_document.doc_form + index_processor = IndexProcessorFactory(index_type).init_index_processor() + # extract + text_docs = self._extract(index_processor, dataset_document, processing_rule.to_dict()) + + # transform + documents = self._transform( + index_processor, dataset, text_docs, dataset_document.doc_language, processing_rule.to_dict() + ) + # save segment + self._load_segments(dataset, dataset_document, documents) + + # load + self._load( + index_processor=index_processor, dataset=dataset, dataset_document=dataset_document, documents=documents + ) + except DocumentIsPausedError: + raise DocumentIsPausedError("Document paused, document id: {}".format(dataset_document.id)) + except ProviderTokenNotInitError as e: + dataset_document.indexing_status = "error" + dataset_document.error = str(e.description) + dataset_document.stopped_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None) + db.session.commit() + except Exception as e: + logging.exception("consume document failed") + dataset_document.indexing_status = "error" + dataset_document.error = str(e) + dataset_document.stopped_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None) + db.session.commit() + + def run_in_indexing_status(self, dataset_document: DatasetDocument): + """Run the indexing process when the index_status is indexing.""" + try: + # get dataset + dataset = Dataset.query.filter_by(id=dataset_document.dataset_id).first() + + if not dataset: + raise ValueError("no dataset found") + + # get exist document_segment list and delete + document_segments = DocumentSegment.query.filter_by( + dataset_id=dataset.id, document_id=dataset_document.id + ).all() + + documents = [] + if document_segments: + for document_segment in document_segments: + # transform segment to node + if document_segment.status != "completed": + document = Document( + page_content=document_segment.content, + metadata={ + "doc_id": document_segment.index_node_id, + "doc_hash": document_segment.index_node_hash, + "document_id": document_segment.document_id, + "dataset_id": document_segment.dataset_id, + }, + ) + if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX: + child_chunks = document_segment.child_chunks + if child_chunks: + child_documents = [] + for child_chunk in child_chunks: + child_document = ChildDocument( + page_content=child_chunk.content, + metadata={ + "doc_id": child_chunk.index_node_id, + "doc_hash": child_chunk.index_node_hash, + "document_id": document_segment.document_id, + "dataset_id": document_segment.dataset_id, + }, + ) + child_documents.append(child_document) + document.children = child_documents + documents.append(document) + + # build index + # get the process rule + processing_rule = ( + db.session.query(DatasetProcessRule) + .filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id) + .first() + ) + + index_type = dataset_document.doc_form + index_processor = IndexProcessorFactory(index_type).init_index_processor() + self._load( + index_processor=index_processor, dataset=dataset, dataset_document=dataset_document, documents=documents + ) + except DocumentIsPausedError: + raise DocumentIsPausedError("Document paused, document id: {}".format(dataset_document.id)) + except ProviderTokenNotInitError as e: + dataset_document.indexing_status = "error" + dataset_document.error = str(e.description) + dataset_document.stopped_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None) + db.session.commit() + except Exception as e: + logging.exception("consume document failed") + dataset_document.indexing_status = "error" + dataset_document.error = str(e) + dataset_document.stopped_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None) + db.session.commit() + + def indexing_estimate( + self, + tenant_id: str, + extract_settings: list[ExtractSetting], + tmp_processing_rule: dict, + doc_form: Optional[str] = None, + doc_language: str = "English", + dataset_id: Optional[str] = None, + indexing_technique: str = "economy", + ) -> IndexingEstimate: + """ + Estimate the indexing for the document. + """ + # check document limit + features = FeatureService.get_features(tenant_id) + if features.billing.enabled: + count = len(extract_settings) + batch_upload_limit = dify_config.BATCH_UPLOAD_LIMIT + if count > batch_upload_limit: + raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.") + + embedding_model_instance = None + if dataset_id: + dataset = Dataset.query.filter_by(id=dataset_id).first() + if not dataset: + raise ValueError("Dataset not found.") + if dataset.indexing_technique == "high_quality" or indexing_technique == "high_quality": + if dataset.embedding_model_provider: + embedding_model_instance = self.model_manager.get_model_instance( + tenant_id=tenant_id, + provider=dataset.embedding_model_provider, + model_type=ModelType.TEXT_EMBEDDING, + model=dataset.embedding_model, + ) + else: + embedding_model_instance = self.model_manager.get_default_model_instance( + tenant_id=tenant_id, + model_type=ModelType.TEXT_EMBEDDING, + ) + else: + if indexing_technique == "high_quality": + embedding_model_instance = self.model_manager.get_default_model_instance( + tenant_id=tenant_id, + model_type=ModelType.TEXT_EMBEDDING, + ) + preview_texts = [] # type: ignore + + total_segments = 0 + index_type = doc_form + index_processor = IndexProcessorFactory(index_type).init_index_processor() + for extract_setting in extract_settings: + # extract + processing_rule = DatasetProcessRule( + mode=tmp_processing_rule["mode"], rules=json.dumps(tmp_processing_rule["rules"]) + ) + text_docs = index_processor.extract(extract_setting, process_rule_mode=tmp_processing_rule["mode"]) + documents = index_processor.transform( + text_docs, + embedding_model_instance=embedding_model_instance, + process_rule=processing_rule.to_dict(), + tenant_id=current_user.current_tenant_id, + doc_language=doc_language, + preview=True, + ) + total_segments += len(documents) + for document in documents: + if len(preview_texts) < 10: + if doc_form and doc_form == "qa_model": + preview_detail = QAPreviewDetail( + question=document.page_content, answer=document.metadata.get("answer") or "" + ) + preview_texts.append(preview_detail) + else: + preview_detail = PreviewDetail(content=document.page_content) # type: ignore + if document.children: + preview_detail.child_chunks = [child.page_content for child in document.children] # type: ignore + preview_texts.append(preview_detail) + + # delete image files and related db records + image_upload_file_ids = get_image_upload_file_ids(document.page_content) + for upload_file_id in image_upload_file_ids: + image_file = db.session.query(UploadFile).filter(UploadFile.id == upload_file_id).first() + try: + if image_file: + storage.delete(image_file.key) + except Exception: + logging.exception( + "Delete image_files failed while indexing_estimate, \ + image_upload_file_is: {}".format(upload_file_id) + ) + db.session.delete(image_file) + + if doc_form and doc_form == "qa_model": + return IndexingEstimate(total_segments=total_segments * 20, qa_preview=preview_texts, preview=[]) + return IndexingEstimate(total_segments=total_segments, preview=preview_texts) # type: ignore + + def _extract( + self, index_processor: BaseIndexProcessor, dataset_document: DatasetDocument, process_rule: dict + ) -> list[Document]: + # load file + if dataset_document.data_source_type not in {"upload_file", "notion_import", "website_crawl"}: + return [] + + data_source_info = dataset_document.data_source_info_dict + text_docs = [] + if dataset_document.data_source_type == "upload_file": + if not data_source_info or "upload_file_id" not in data_source_info: + raise ValueError("no upload file found") + + file_detail = ( + db.session.query(UploadFile).filter(UploadFile.id == data_source_info["upload_file_id"]).one_or_none() + ) + + if file_detail: + extract_setting = ExtractSetting( + datasource_type="upload_file", upload_file=file_detail, document_model=dataset_document.doc_form + ) + text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule["mode"]) + elif dataset_document.data_source_type == "notion_import": + if ( + not data_source_info + or "notion_workspace_id" not in data_source_info + or "notion_page_id" not in data_source_info + ): + raise ValueError("no notion import info found") + extract_setting = ExtractSetting( + datasource_type="notion_import", + notion_info={ + "notion_workspace_id": data_source_info["notion_workspace_id"], + "notion_obj_id": data_source_info["notion_page_id"], + "notion_page_type": data_source_info["type"], + "document": dataset_document, + "tenant_id": dataset_document.tenant_id, + }, + document_model=dataset_document.doc_form, + ) + text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule["mode"]) + elif dataset_document.data_source_type == "website_crawl": + if ( + not data_source_info + or "provider" not in data_source_info + or "url" not in data_source_info + or "job_id" not in data_source_info + ): + raise ValueError("no website import info found") + extract_setting = ExtractSetting( + datasource_type="website_crawl", + website_info={ + "provider": data_source_info["provider"], + "job_id": data_source_info["job_id"], + "tenant_id": dataset_document.tenant_id, + "url": data_source_info["url"], + "mode": data_source_info["mode"], + "only_main_content": data_source_info["only_main_content"], + }, + document_model=dataset_document.doc_form, + ) + text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule["mode"]) + # update document status to splitting + self._update_document_index_status( + document_id=dataset_document.id, + after_indexing_status="splitting", + extra_update_params={ + DatasetDocument.word_count: sum(len(text_doc.page_content) for text_doc in text_docs), + DatasetDocument.parsing_completed_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None), + }, + ) + + # replace doc id to document model id + text_docs = cast(list[Document], text_docs) + for text_doc in text_docs: + if text_doc.metadata is not None: + text_doc.metadata["document_id"] = dataset_document.id + text_doc.metadata["dataset_id"] = dataset_document.dataset_id + + return text_docs + + @staticmethod + def filter_string(text): + text = re.sub(r"<\|", "<", text) + text = re.sub(r"\|>", ">", text) + text = re.sub(r"[\x00-\x08\x0B\x0C\x0E-\x1F\x7F\xEF\xBF\xBE]", "", text) + # Unicode U+FFFE + text = re.sub("\ufffe", "", text) + return text + + @staticmethod + def _get_splitter( + processing_rule_mode: str, + max_tokens: int, + chunk_overlap: int, + separator: str, + embedding_model_instance: Optional[ModelInstance], + ) -> TextSplitter: + """ + Get the NodeParser object according to the processing rule. + """ + if processing_rule_mode in ["custom", "hierarchical"]: + # The user-defined segmentation rule + max_segmentation_tokens_length = dify_config.INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH + if max_tokens < 50 or max_tokens > max_segmentation_tokens_length: + raise ValueError(f"Custom segment length should be between 50 and {max_segmentation_tokens_length}.") + + if separator: + separator = separator.replace("\\n", "\n") + + character_splitter = FixedRecursiveCharacterTextSplitter.from_encoder( + chunk_size=max_tokens, + chunk_overlap=chunk_overlap, + fixed_separator=separator, + separators=["\n\n", "。", ". ", " ", ""], + embedding_model_instance=embedding_model_instance, + ) + else: + # Automatic segmentation + automatic_rules: dict[str, Any] = dict(DatasetProcessRule.AUTOMATIC_RULES["segmentation"]) + character_splitter = EnhanceRecursiveCharacterTextSplitter.from_encoder( + chunk_size=automatic_rules["max_tokens"], + chunk_overlap=automatic_rules["chunk_overlap"], + separators=["\n\n", "。", ". ", " ", ""], + embedding_model_instance=embedding_model_instance, + ) + + return character_splitter # type: ignore + + def _split_to_documents_for_estimate( + self, text_docs: list[Document], splitter: TextSplitter, processing_rule: DatasetProcessRule + ) -> list[Document]: + """ + Split the text documents into nodes. + """ + all_documents: list[Document] = [] + for text_doc in text_docs: + # document clean + document_text = self._document_clean(text_doc.page_content, processing_rule) + text_doc.page_content = document_text + + # parse document to nodes + documents = splitter.split_documents([text_doc]) + + split_documents = [] + for document in documents: + if document.page_content is None or not document.page_content.strip(): + continue + if document.metadata is not None: + doc_id = str(uuid.uuid4()) + hash = helper.generate_text_hash(document.page_content) + document.metadata["doc_id"] = doc_id + document.metadata["doc_hash"] = hash + + split_documents.append(document) + + all_documents.extend(split_documents) + + return all_documents + + @staticmethod + def _document_clean(text: str, processing_rule: DatasetProcessRule) -> str: + """ + Clean the document text according to the processing rules. + """ + if processing_rule.mode == "automatic": + rules = DatasetProcessRule.AUTOMATIC_RULES + else: + rules = json.loads(processing_rule.rules) if processing_rule.rules else {} + document_text = CleanProcessor.clean(text, {"rules": rules}) + + return document_text + + @staticmethod + def format_split_text(text: str) -> list[QAPreviewDetail]: + regex = r"Q\d+:\s*(.*?)\s*A\d+:\s*([\s\S]*?)(?=Q\d+:|$)" + matches = re.findall(regex, text, re.UNICODE) + + return [QAPreviewDetail(question=q, answer=re.sub(r"\n\s*", "\n", a.strip())) for q, a in matches if q and a] + + def _load( + self, + index_processor: BaseIndexProcessor, + dataset: Dataset, + dataset_document: DatasetDocument, + documents: list[Document], + ) -> None: + """ + insert index and update document/segment status to completed + """ + + embedding_model_instance = None + if dataset.indexing_technique == "high_quality": + embedding_model_instance = self.model_manager.get_model_instance( + tenant_id=dataset.tenant_id, + provider=dataset.embedding_model_provider, + model_type=ModelType.TEXT_EMBEDDING, + model=dataset.embedding_model, + ) + + # chunk nodes by chunk size + indexing_start_at = time.perf_counter() + tokens = 0 + if dataset_document.doc_form != IndexType.PARENT_CHILD_INDEX: + # create keyword index + create_keyword_thread = threading.Thread( + target=self._process_keyword_index, + args=(current_app._get_current_object(), dataset.id, dataset_document.id, documents), # type: ignore + ) + create_keyword_thread.start() + + max_workers = 10 + if dataset.indexing_technique == "high_quality": + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [] + + # Distribute documents into multiple groups based on the hash values of page_content + # This is done to prevent multiple threads from processing the same document, + # Thereby avoiding potential database insertion deadlocks + document_groups: list[list[Document]] = [[] for _ in range(max_workers)] + for document in documents: + hash = helper.generate_text_hash(document.page_content) + group_index = int(hash, 16) % max_workers + document_groups[group_index].append(document) + for chunk_documents in document_groups: + if len(chunk_documents) == 0: + continue + futures.append( + executor.submit( + self._process_chunk, + current_app._get_current_object(), # type: ignore + index_processor, + chunk_documents, + dataset, + dataset_document, + embedding_model_instance, + ) + ) + + for future in futures: + tokens += future.result() + if dataset_document.doc_form != IndexType.PARENT_CHILD_INDEX: + create_keyword_thread.join() + indexing_end_at = time.perf_counter() + + # update document status to completed + self._update_document_index_status( + document_id=dataset_document.id, + after_indexing_status="completed", + extra_update_params={ + DatasetDocument.tokens: tokens, + DatasetDocument.completed_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None), + DatasetDocument.indexing_latency: indexing_end_at - indexing_start_at, + DatasetDocument.error: None, + }, + ) + + @staticmethod + def _process_keyword_index(flask_app, dataset_id, document_id, documents): + with flask_app.app_context(): + dataset = Dataset.query.filter_by(id=dataset_id).first() + if not dataset: + raise ValueError("no dataset found") + keyword = Keyword(dataset) + keyword.create(documents) + if dataset.indexing_technique != "high_quality": + document_ids = [document.metadata["doc_id"] for document in documents] + db.session.query(DocumentSegment).filter( + DocumentSegment.document_id == document_id, + DocumentSegment.dataset_id == dataset_id, + DocumentSegment.index_node_id.in_(document_ids), + DocumentSegment.status == "indexing", + ).update( + { + DocumentSegment.status: "completed", + DocumentSegment.enabled: True, + DocumentSegment.completed_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None), + } + ) + + db.session.commit() + + def _process_chunk( + self, flask_app, index_processor, chunk_documents, dataset, dataset_document, embedding_model_instance + ): + with flask_app.app_context(): + # check document is paused + self._check_document_paused_status(dataset_document.id) + + tokens = 0 + if embedding_model_instance: + tokens += sum( + embedding_model_instance.get_text_embedding_num_tokens([document.page_content]) + for document in chunk_documents + ) + + # load index + index_processor.load(dataset, chunk_documents, with_keywords=False) + + document_ids = [document.metadata["doc_id"] for document in chunk_documents] + db.session.query(DocumentSegment).filter( + DocumentSegment.document_id == dataset_document.id, + DocumentSegment.dataset_id == dataset.id, + DocumentSegment.index_node_id.in_(document_ids), + DocumentSegment.status == "indexing", + ).update( + { + DocumentSegment.status: "completed", + DocumentSegment.enabled: True, + DocumentSegment.completed_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None), + } + ) + + db.session.commit() + + return tokens + + @staticmethod + def _check_document_paused_status(document_id: str): + indexing_cache_key = "document_{}_is_paused".format(document_id) + result = redis_client.get(indexing_cache_key) + if result: + raise DocumentIsPausedError() + + @staticmethod + def _update_document_index_status( + document_id: str, after_indexing_status: str, extra_update_params: Optional[dict] = None + ) -> None: + """ + Update the document indexing status. + """ + count = DatasetDocument.query.filter_by(id=document_id, is_paused=True).count() + if count > 0: + raise DocumentIsPausedError() + document = DatasetDocument.query.filter_by(id=document_id).first() + if not document: + raise DocumentIsDeletedPausedError() + + update_params = {DatasetDocument.indexing_status: after_indexing_status} + + if extra_update_params: + update_params.update(extra_update_params) + + DatasetDocument.query.filter_by(id=document_id).update(update_params) + db.session.commit() + + @staticmethod + def _update_segments_by_document(dataset_document_id: str, update_params: dict) -> None: + """ + Update the document segment by document id. + """ + DocumentSegment.query.filter_by(document_id=dataset_document_id).update(update_params) + db.session.commit() + + def _transform( + self, + index_processor: BaseIndexProcessor, + dataset: Dataset, + text_docs: list[Document], + doc_language: str, + process_rule: dict, + ) -> list[Document]: + # get embedding model instance + embedding_model_instance = None + if dataset.indexing_technique == "high_quality": + if dataset.embedding_model_provider: + embedding_model_instance = self.model_manager.get_model_instance( + tenant_id=dataset.tenant_id, + provider=dataset.embedding_model_provider, + model_type=ModelType.TEXT_EMBEDDING, + model=dataset.embedding_model, + ) + else: + embedding_model_instance = self.model_manager.get_default_model_instance( + tenant_id=dataset.tenant_id, + model_type=ModelType.TEXT_EMBEDDING, + ) + + documents = index_processor.transform( + text_docs, + embedding_model_instance=embedding_model_instance, + process_rule=process_rule, + tenant_id=dataset.tenant_id, + doc_language=doc_language, + ) + + return documents + + def _load_segments(self, dataset, dataset_document, documents): + # save node to document segment + doc_store = DatasetDocumentStore( + dataset=dataset, user_id=dataset_document.created_by, document_id=dataset_document.id + ) + + # add document segments + doc_store.add_documents(docs=documents, save_child=dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX) + + # update document status to indexing + cur_time = datetime.datetime.now(datetime.UTC).replace(tzinfo=None) + self._update_document_index_status( + document_id=dataset_document.id, + after_indexing_status="indexing", + extra_update_params={ + DatasetDocument.cleaning_completed_at: cur_time, + DatasetDocument.splitting_completed_at: cur_time, + }, + ) + + # update segment status to indexing + self._update_segments_by_document( + dataset_document_id=dataset_document.id, + update_params={ + DocumentSegment.status: "indexing", + DocumentSegment.indexing_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None), + }, + ) + pass + + +class DocumentIsPausedError(Exception): + pass + + +class DocumentIsDeletedPausedError(Exception): + pass diff --git a/api/core/model_manager.py b/api/core/model_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..d1e71148cd6023b194fffa28c2feb63f93bc6dd3 --- /dev/null +++ b/api/core/model_manager.py @@ -0,0 +1,559 @@ +import logging +from collections.abc import Callable, Generator, Iterable, Sequence +from typing import IO, Any, Optional, Union, cast + +from configs import dify_config +from core.entities.embedding_type import EmbeddingInputType +from core.entities.provider_configuration import ProviderConfiguration, ProviderModelBundle +from core.entities.provider_entities import ModelLoadBalancingConfiguration +from core.errors.error import ProviderTokenNotInitError +from core.model_runtime.callbacks.base_callback import Callback +from core.model_runtime.entities.llm_entities import LLMResult +from core.model_runtime.entities.message_entities import PromptMessage, PromptMessageTool +from core.model_runtime.entities.model_entities import ModelType +from core.model_runtime.entities.rerank_entities import RerankResult +from core.model_runtime.entities.text_embedding_entities import TextEmbeddingResult +from core.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeConnectionError, InvokeRateLimitError +from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel +from core.model_runtime.model_providers.__base.moderation_model import ModerationModel +from core.model_runtime.model_providers.__base.rerank_model import RerankModel +from core.model_runtime.model_providers.__base.speech2text_model import Speech2TextModel +from core.model_runtime.model_providers.__base.text_embedding_model import TextEmbeddingModel +from core.model_runtime.model_providers.__base.tts_model import TTSModel +from core.provider_manager import ProviderManager +from extensions.ext_redis import redis_client +from models.provider import ProviderType + +logger = logging.getLogger(__name__) + + +class ModelInstance: + """ + Model instance class + """ + + def __init__(self, provider_model_bundle: ProviderModelBundle, model: str) -> None: + self.provider_model_bundle = provider_model_bundle + self.model = model + self.provider = provider_model_bundle.configuration.provider.provider + self.credentials = self._fetch_credentials_from_bundle(provider_model_bundle, model) + self.model_type_instance = self.provider_model_bundle.model_type_instance + self.load_balancing_manager = self._get_load_balancing_manager( + configuration=provider_model_bundle.configuration, + model_type=provider_model_bundle.model_type_instance.model_type, + model=model, + credentials=self.credentials, + ) + + @staticmethod + def _fetch_credentials_from_bundle(provider_model_bundle: ProviderModelBundle, model: str) -> dict: + """ + Fetch credentials from provider model bundle + :param provider_model_bundle: provider model bundle + :param model: model name + :return: + """ + configuration = provider_model_bundle.configuration + model_type = provider_model_bundle.model_type_instance.model_type + credentials = configuration.get_current_credentials(model_type=model_type, model=model) + + if credentials is None: + raise ProviderTokenNotInitError(f"Model {model} credentials is not initialized.") + + return credentials + + @staticmethod + def _get_load_balancing_manager( + configuration: ProviderConfiguration, model_type: ModelType, model: str, credentials: dict + ) -> Optional["LBModelManager"]: + """ + Get load balancing model credentials + :param configuration: provider configuration + :param model_type: model type + :param model: model name + :param credentials: model credentials + :return: + """ + if configuration.model_settings and configuration.using_provider_type == ProviderType.CUSTOM: + current_model_setting = None + # check if model is disabled by admin + for model_setting in configuration.model_settings: + if model_setting.model_type == model_type and model_setting.model == model: + current_model_setting = model_setting + break + + # check if load balancing is enabled + if current_model_setting and current_model_setting.load_balancing_configs: + # use load balancing proxy to choose credentials + lb_model_manager = LBModelManager( + tenant_id=configuration.tenant_id, + provider=configuration.provider.provider, + model_type=model_type, + model=model, + load_balancing_configs=current_model_setting.load_balancing_configs, + managed_credentials=credentials if configuration.custom_configuration.provider else None, + ) + + return lb_model_manager + + return None + + def invoke_llm( + self, + prompt_messages: Sequence[PromptMessage], + model_parameters: Optional[dict] = None, + tools: Sequence[PromptMessageTool] | None = None, + stop: Optional[Sequence[str]] = None, + stream: bool = True, + user: Optional[str] = None, + callbacks: Optional[list[Callback]] = None, + ) -> Union[LLMResult, Generator]: + """ + Invoke large language model + + :param prompt_messages: prompt messages + :param model_parameters: model parameters + :param tools: tools for tool calling + :param stop: stop words + :param stream: is stream response + :param user: unique user id + :param callbacks: callbacks + :return: full response or stream response chunk generator result + """ + if not isinstance(self.model_type_instance, LargeLanguageModel): + raise Exception("Model type instance is not LargeLanguageModel") + + self.model_type_instance = cast(LargeLanguageModel, self.model_type_instance) + return cast( + Union[LLMResult, Generator], + self._round_robin_invoke( + function=self.model_type_instance.invoke, + model=self.model, + credentials=self.credentials, + prompt_messages=prompt_messages, + model_parameters=model_parameters, + tools=tools, + stop=stop, + stream=stream, + user=user, + callbacks=callbacks, + ), + ) + + def get_llm_num_tokens( + self, prompt_messages: list[PromptMessage], tools: Optional[list[PromptMessageTool]] = None + ) -> int: + """ + Get number of tokens for llm + + :param prompt_messages: prompt messages + :param tools: tools for tool calling + :return: + """ + if not isinstance(self.model_type_instance, LargeLanguageModel): + raise Exception("Model type instance is not LargeLanguageModel") + + self.model_type_instance = cast(LargeLanguageModel, self.model_type_instance) + return cast( + int, + self._round_robin_invoke( + function=self.model_type_instance.get_num_tokens, + model=self.model, + credentials=self.credentials, + prompt_messages=prompt_messages, + tools=tools, + ), + ) + + def invoke_text_embedding( + self, texts: list[str], user: Optional[str] = None, input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT + ) -> TextEmbeddingResult: + """ + Invoke large language model + + :param texts: texts to embed + :param user: unique user id + :param input_type: input type + :return: embeddings result + """ + if not isinstance(self.model_type_instance, TextEmbeddingModel): + raise Exception("Model type instance is not TextEmbeddingModel") + + self.model_type_instance = cast(TextEmbeddingModel, self.model_type_instance) + return cast( + TextEmbeddingResult, + self._round_robin_invoke( + function=self.model_type_instance.invoke, + model=self.model, + credentials=self.credentials, + texts=texts, + user=user, + input_type=input_type, + ), + ) + + def get_text_embedding_num_tokens(self, texts: list[str]) -> int: + """ + Get number of tokens for text embedding + + :param texts: texts to embed + :return: + """ + if not isinstance(self.model_type_instance, TextEmbeddingModel): + raise Exception("Model type instance is not TextEmbeddingModel") + + self.model_type_instance = cast(TextEmbeddingModel, self.model_type_instance) + return cast( + int, + self._round_robin_invoke( + function=self.model_type_instance.get_num_tokens, + model=self.model, + credentials=self.credentials, + texts=texts, + ), + ) + + def invoke_rerank( + self, + query: str, + docs: list[str], + score_threshold: Optional[float] = None, + top_n: Optional[int] = None, + user: Optional[str] = None, + ) -> RerankResult: + """ + Invoke rerank model + + :param query: search query + :param docs: docs for reranking + :param score_threshold: score threshold + :param top_n: top n + :param user: unique user id + :return: rerank result + """ + if not isinstance(self.model_type_instance, RerankModel): + raise Exception("Model type instance is not RerankModel") + + self.model_type_instance = cast(RerankModel, self.model_type_instance) + return cast( + RerankResult, + self._round_robin_invoke( + function=self.model_type_instance.invoke, + model=self.model, + credentials=self.credentials, + query=query, + docs=docs, + score_threshold=score_threshold, + top_n=top_n, + user=user, + ), + ) + + def invoke_moderation(self, text: str, user: Optional[str] = None) -> bool: + """ + Invoke moderation model + + :param text: text to moderate + :param user: unique user id + :return: false if text is safe, true otherwise + """ + if not isinstance(self.model_type_instance, ModerationModel): + raise Exception("Model type instance is not ModerationModel") + + self.model_type_instance = cast(ModerationModel, self.model_type_instance) + return cast( + bool, + self._round_robin_invoke( + function=self.model_type_instance.invoke, + model=self.model, + credentials=self.credentials, + text=text, + user=user, + ), + ) + + def invoke_speech2text(self, file: IO[bytes], user: Optional[str] = None) -> str: + """ + Invoke large language model + + :param file: audio file + :param user: unique user id + :return: text for given audio file + """ + if not isinstance(self.model_type_instance, Speech2TextModel): + raise Exception("Model type instance is not Speech2TextModel") + + self.model_type_instance = cast(Speech2TextModel, self.model_type_instance) + return cast( + str, + self._round_robin_invoke( + function=self.model_type_instance.invoke, + model=self.model, + credentials=self.credentials, + file=file, + user=user, + ), + ) + + def invoke_tts(self, content_text: str, tenant_id: str, voice: str, user: Optional[str] = None) -> Iterable[bytes]: + """ + Invoke large language tts model + + :param content_text: text content to be translated + :param tenant_id: user tenant id + :param voice: model timbre + :param user: unique user id + :return: text for given audio file + """ + if not isinstance(self.model_type_instance, TTSModel): + raise Exception("Model type instance is not TTSModel") + + self.model_type_instance = cast(TTSModel, self.model_type_instance) + return cast( + Iterable[bytes], + self._round_robin_invoke( + function=self.model_type_instance.invoke, + model=self.model, + credentials=self.credentials, + content_text=content_text, + user=user, + tenant_id=tenant_id, + voice=voice, + ), + ) + + def _round_robin_invoke(self, function: Callable[..., Any], *args, **kwargs) -> Any: + """ + Round-robin invoke + :param function: function to invoke + :param args: function args + :param kwargs: function kwargs + :return: + """ + if not self.load_balancing_manager: + return function(*args, **kwargs) + + last_exception: Union[InvokeRateLimitError, InvokeAuthorizationError, InvokeConnectionError, None] = None + while True: + lb_config = self.load_balancing_manager.fetch_next() + if not lb_config: + if not last_exception: + raise ProviderTokenNotInitError("Model credentials is not initialized.") + else: + raise last_exception + + try: + if "credentials" in kwargs: + del kwargs["credentials"] + return function(*args, **kwargs, credentials=lb_config.credentials) + except InvokeRateLimitError as e: + # expire in 60 seconds + self.load_balancing_manager.cooldown(lb_config, expire=60) + last_exception = e + continue + except (InvokeAuthorizationError, InvokeConnectionError) as e: + # expire in 10 seconds + self.load_balancing_manager.cooldown(lb_config, expire=10) + last_exception = e + continue + except Exception as e: + raise e + + def get_tts_voices(self, language: Optional[str] = None) -> list: + """ + Invoke large language tts model voices + + :param language: tts language + :return: tts model voices + """ + if not isinstance(self.model_type_instance, TTSModel): + raise Exception("Model type instance is not TTSModel") + + self.model_type_instance = cast(TTSModel, self.model_type_instance) + return self.model_type_instance.get_tts_model_voices( + model=self.model, credentials=self.credentials, language=language + ) + + +class ModelManager: + def __init__(self) -> None: + self._provider_manager = ProviderManager() + + def get_model_instance(self, tenant_id: str, provider: str, model_type: ModelType, model: str) -> ModelInstance: + """ + Get model instance + :param tenant_id: tenant id + :param provider: provider name + :param model_type: model type + :param model: model name + :return: + """ + if not provider: + return self.get_default_model_instance(tenant_id, model_type) + + provider_model_bundle = self._provider_manager.get_provider_model_bundle( + tenant_id=tenant_id, provider=provider, model_type=model_type + ) + + return ModelInstance(provider_model_bundle, model) + + def get_default_provider_model_name(self, tenant_id: str, model_type: ModelType) -> tuple[str, str]: + """ + Return first provider and the first model in the provider + :param tenant_id: tenant id + :param model_type: model type + :return: provider name, model name + """ + return self._provider_manager.get_first_provider_first_model(tenant_id, model_type) + + def get_default_model_instance(self, tenant_id: str, model_type: ModelType) -> ModelInstance: + """ + Get default model instance + :param tenant_id: tenant id + :param model_type: model type + :return: + """ + default_model_entity = self._provider_manager.get_default_model(tenant_id=tenant_id, model_type=model_type) + + if not default_model_entity: + raise ProviderTokenNotInitError(f"Default model not found for {model_type}") + + return self.get_model_instance( + tenant_id=tenant_id, + provider=default_model_entity.provider.provider, + model_type=model_type, + model=default_model_entity.model, + ) + + +class LBModelManager: + def __init__( + self, + tenant_id: str, + provider: str, + model_type: ModelType, + model: str, + load_balancing_configs: list[ModelLoadBalancingConfiguration], + managed_credentials: Optional[dict] = None, + ) -> None: + """ + Load balancing model manager + :param tenant_id: tenant_id + :param provider: provider + :param model_type: model_type + :param model: model name + :param load_balancing_configs: all load balancing configurations + :param managed_credentials: credentials if load balancing configuration name is __inherit__ + """ + self._tenant_id = tenant_id + self._provider = provider + self._model_type = model_type + self._model = model + self._load_balancing_configs = load_balancing_configs + + for load_balancing_config in self._load_balancing_configs[:]: # Iterate over a shallow copy of the list + if load_balancing_config.name == "__inherit__": + if not managed_credentials: + # remove __inherit__ if managed credentials is not provided + self._load_balancing_configs.remove(load_balancing_config) + else: + load_balancing_config.credentials = managed_credentials + + def fetch_next(self) -> Optional[ModelLoadBalancingConfiguration]: + """ + Get next model load balancing config + Strategy: Round Robin + :return: + """ + cache_key = "model_lb_index:{}:{}:{}:{}".format( + self._tenant_id, self._provider, self._model_type.value, self._model + ) + + cooldown_load_balancing_configs = [] + max_index = len(self._load_balancing_configs) + + while True: + current_index = redis_client.incr(cache_key) + current_index = cast(int, current_index) + if current_index >= 10000000: + current_index = 1 + redis_client.set(cache_key, current_index) + + redis_client.expire(cache_key, 3600) + if current_index > max_index: + current_index = current_index % max_index + + real_index = current_index - 1 + if real_index > max_index: + real_index = 0 + + config: ModelLoadBalancingConfiguration = self._load_balancing_configs[real_index] + + if self.in_cooldown(config): + cooldown_load_balancing_configs.append(config) + if len(cooldown_load_balancing_configs) >= len(self._load_balancing_configs): + # all configs are in cooldown + return None + + continue + + if dify_config.DEBUG: + logger.info( + f"Model LB\nid: {config.id}\nname:{config.name}\n" + f"tenant_id: {self._tenant_id}\nprovider: {self._provider}\n" + f"model_type: {self._model_type.value}\nmodel: {self._model}" + ) + + return config + + return None + + def cooldown(self, config: ModelLoadBalancingConfiguration, expire: int = 60) -> None: + """ + Cooldown model load balancing config + :param config: model load balancing config + :param expire: cooldown time + :return: + """ + cooldown_cache_key = "model_lb_index:cooldown:{}:{}:{}:{}:{}".format( + self._tenant_id, self._provider, self._model_type.value, self._model, config.id + ) + + redis_client.setex(cooldown_cache_key, expire, "true") + + def in_cooldown(self, config: ModelLoadBalancingConfiguration) -> bool: + """ + Check if model load balancing config is in cooldown + :param config: model load balancing config + :return: + """ + cooldown_cache_key = "model_lb_index:cooldown:{}:{}:{}:{}:{}".format( + self._tenant_id, self._provider, self._model_type.value, self._model, config.id + ) + + res: bool = redis_client.exists(cooldown_cache_key) + return res + + @staticmethod + def get_config_in_cooldown_and_ttl( + tenant_id: str, provider: str, model_type: ModelType, model: str, config_id: str + ) -> tuple[bool, int]: + """ + Get model load balancing config is in cooldown and ttl + :param tenant_id: workspace id + :param provider: provider name + :param model_type: model type + :param model: model name + :param config_id: model load balancing config id + :return: + """ + cooldown_cache_key = "model_lb_index:cooldown:{}:{}:{}:{}:{}".format( + tenant_id, provider, model_type.value, model, config_id + ) + + ttl = redis_client.ttl(cooldown_cache_key) + if ttl == -2: + return False, 0 + + ttl = cast(int, ttl) + return True, ttl diff --git a/api/core/moderation/__init__.py b/api/core/moderation/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/moderation/api/__builtin__ b/api/core/moderation/api/__builtin__ new file mode 100644 index 0000000000000000000000000000000000000000..e440e5c842586965a7fb77deda2eca68612b1f53 --- /dev/null +++ b/api/core/moderation/api/__builtin__ @@ -0,0 +1 @@ +3 \ No newline at end of file diff --git a/api/core/moderation/api/__init__.py b/api/core/moderation/api/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/moderation/api/api.py b/api/core/moderation/api/api.py new file mode 100644 index 0000000000000000000000000000000000000000..c65a3885fd1eb99673364705dfa96f9d88cbc7d2 --- /dev/null +++ b/api/core/moderation/api/api.py @@ -0,0 +1,96 @@ +from typing import Optional + +from pydantic import BaseModel + +from core.extension.api_based_extension_requestor import APIBasedExtensionPoint, APIBasedExtensionRequestor +from core.helper.encrypter import decrypt_token +from core.moderation.base import Moderation, ModerationAction, ModerationInputsResult, ModerationOutputsResult +from extensions.ext_database import db +from models.api_based_extension import APIBasedExtension + + +class ModerationInputParams(BaseModel): + app_id: str = "" + inputs: dict = {} + query: str = "" + + +class ModerationOutputParams(BaseModel): + app_id: str = "" + text: str + + +class ApiModeration(Moderation): + name: str = "api" + + @classmethod + def validate_config(cls, tenant_id: str, config: dict) -> None: + """ + Validate the incoming form config data. + + :param tenant_id: the id of workspace + :param config: the form config data + :return: + """ + cls._validate_inputs_and_outputs_config(config, False) + + api_based_extension_id = config.get("api_based_extension_id") + if not api_based_extension_id: + raise ValueError("api_based_extension_id is required") + + extension = cls._get_api_based_extension(tenant_id, api_based_extension_id) + if not extension: + raise ValueError("API-based Extension not found. Please check it again.") + + def moderation_for_inputs(self, inputs: dict, query: str = "") -> ModerationInputsResult: + flagged = False + preset_response = "" + if self.config is None: + raise ValueError("The config is not set.") + + if self.config["inputs_config"]["enabled"]: + params = ModerationInputParams(app_id=self.app_id, inputs=inputs, query=query) + + result = self._get_config_by_requestor(APIBasedExtensionPoint.APP_MODERATION_INPUT, params.model_dump()) + return ModerationInputsResult(**result) + + return ModerationInputsResult( + flagged=flagged, action=ModerationAction.DIRECT_OUTPUT, preset_response=preset_response + ) + + def moderation_for_outputs(self, text: str) -> ModerationOutputsResult: + flagged = False + preset_response = "" + if self.config is None: + raise ValueError("The config is not set.") + + if self.config["outputs_config"]["enabled"]: + params = ModerationOutputParams(app_id=self.app_id, text=text) + + result = self._get_config_by_requestor(APIBasedExtensionPoint.APP_MODERATION_OUTPUT, params.model_dump()) + return ModerationOutputsResult(**result) + + return ModerationOutputsResult( + flagged=flagged, action=ModerationAction.DIRECT_OUTPUT, preset_response=preset_response + ) + + def _get_config_by_requestor(self, extension_point: APIBasedExtensionPoint, params: dict) -> dict: + if self.config is None: + raise ValueError("The config is not set.") + extension = self._get_api_based_extension(self.tenant_id, self.config.get("api_based_extension_id", "")) + if not extension: + raise ValueError("API-based Extension not found. Please check it again.") + requestor = APIBasedExtensionRequestor(extension.api_endpoint, decrypt_token(self.tenant_id, extension.api_key)) + + result = requestor.request(extension_point, params) + return result + + @staticmethod + def _get_api_based_extension(tenant_id: str, api_based_extension_id: str) -> Optional[APIBasedExtension]: + extension = ( + db.session.query(APIBasedExtension) + .filter(APIBasedExtension.tenant_id == tenant_id, APIBasedExtension.id == api_based_extension_id) + .first() + ) + + return extension diff --git a/api/core/moderation/base.py b/api/core/moderation/base.py new file mode 100644 index 0000000000000000000000000000000000000000..d8c392d0970e1917a986b04efe31ce50031f1b9f --- /dev/null +++ b/api/core/moderation/base.py @@ -0,0 +1,115 @@ +from abc import ABC, abstractmethod +from enum import Enum +from typing import Optional + +from pydantic import BaseModel + +from core.extension.extensible import Extensible, ExtensionModule + + +class ModerationAction(Enum): + DIRECT_OUTPUT = "direct_output" + OVERRIDDEN = "overridden" + + +class ModerationInputsResult(BaseModel): + flagged: bool = False + action: ModerationAction + preset_response: str = "" + inputs: dict = {} + query: str = "" + + +class ModerationOutputsResult(BaseModel): + flagged: bool = False + action: ModerationAction + preset_response: str = "" + text: str = "" + + +class Moderation(Extensible, ABC): + """ + The base class of moderation. + """ + + module: ExtensionModule = ExtensionModule.MODERATION + + def __init__(self, app_id: str, tenant_id: str, config: Optional[dict] = None) -> None: + super().__init__(tenant_id, config) + self.app_id = app_id + + @classmethod + @abstractmethod + def validate_config(cls, tenant_id: str, config: dict) -> None: + """ + Validate the incoming form config data. + + :param tenant_id: the id of workspace + :param config: the form config data + :return: + """ + raise NotImplementedError + + @abstractmethod + def moderation_for_inputs(self, inputs: dict, query: str = "") -> ModerationInputsResult: + """ + Moderation for inputs. + After the user inputs, this method will be called to perform sensitive content review + on the user inputs and return the processed results. + + :param inputs: user inputs + :param query: query string (required in chat app) + :return: + """ + raise NotImplementedError + + @abstractmethod + def moderation_for_outputs(self, text: str) -> ModerationOutputsResult: + """ + Moderation for outputs. + When LLM outputs content, the front end will pass the output content (may be segmented) + to this method for sensitive content review, and the output content will be shielded if the review fails. + + :param text: LLM output content + :return: + """ + raise NotImplementedError + + @classmethod + def _validate_inputs_and_outputs_config(cls, config: dict, is_preset_response_required: bool) -> None: + # inputs_config + inputs_config = config.get("inputs_config") + if not isinstance(inputs_config, dict): + raise ValueError("inputs_config must be a dict") + + # outputs_config + outputs_config = config.get("outputs_config") + if not isinstance(outputs_config, dict): + raise ValueError("outputs_config must be a dict") + + inputs_config_enabled = inputs_config.get("enabled") + outputs_config_enabled = outputs_config.get("enabled") + if not inputs_config_enabled and not outputs_config_enabled: + raise ValueError("At least one of inputs_config or outputs_config must be enabled") + + # preset_response + if not is_preset_response_required: + return + + if inputs_config_enabled: + if not inputs_config.get("preset_response"): + raise ValueError("inputs_config.preset_response is required") + + if len(inputs_config.get("preset_response", 0)) > 100: + raise ValueError("inputs_config.preset_response must be less than 100 characters") + + if outputs_config_enabled: + if not outputs_config.get("preset_response"): + raise ValueError("outputs_config.preset_response is required") + + if len(outputs_config.get("preset_response", 0)) > 100: + raise ValueError("outputs_config.preset_response must be less than 100 characters") + + +class ModerationError(Exception): + pass diff --git a/api/core/moderation/factory.py b/api/core/moderation/factory.py new file mode 100644 index 0000000000000000000000000000000000000000..0ad4438c143870c5d1593aa3ec10a5ade45cffc2 --- /dev/null +++ b/api/core/moderation/factory.py @@ -0,0 +1,49 @@ +from core.extension.extensible import ExtensionModule +from core.moderation.base import Moderation, ModerationInputsResult, ModerationOutputsResult +from extensions.ext_code_based_extension import code_based_extension + + +class ModerationFactory: + __extension_instance: Moderation + + def __init__(self, name: str, app_id: str, tenant_id: str, config: dict) -> None: + extension_class = code_based_extension.extension_class(ExtensionModule.MODERATION, name) + self.__extension_instance = extension_class(app_id, tenant_id, config) + + @classmethod + def validate_config(cls, name: str, tenant_id: str, config: dict) -> None: + """ + Validate the incoming form config data. + + :param name: the name of extension + :param tenant_id: the id of workspace + :param config: the form config data + :return: + """ + code_based_extension.validate_form_schema(ExtensionModule.MODERATION, name, config) + extension_class = code_based_extension.extension_class(ExtensionModule.MODERATION, name) + # FIXME: mypy error, try to fix it instead of using type: ignore + extension_class.validate_config(tenant_id, config) # type: ignore + + def moderation_for_inputs(self, inputs: dict, query: str = "") -> ModerationInputsResult: + """ + Moderation for inputs. + After the user inputs, this method will be called to perform sensitive content review + on the user inputs and return the processed results. + + :param inputs: user inputs + :param query: query string (required in chat app) + :return: + """ + return self.__extension_instance.moderation_for_inputs(inputs, query) + + def moderation_for_outputs(self, text: str) -> ModerationOutputsResult: + """ + Moderation for outputs. + When LLM outputs content, the front end will pass the output content (may be segmented) + to this method for sensitive content review, and the output content will be shielded if the review fails. + + :param text: LLM output content + :return: + """ + return self.__extension_instance.moderation_for_outputs(text) diff --git a/api/core/moderation/input_moderation.py b/api/core/moderation/input_moderation.py new file mode 100644 index 0000000000000000000000000000000000000000..3ac33966cb14bf63c8148ce444037a55e9e55f47 --- /dev/null +++ b/api/core/moderation/input_moderation.py @@ -0,0 +1,71 @@ +import logging +from collections.abc import Mapping +from typing import Any, Optional + +from core.app.app_config.entities import AppConfig +from core.moderation.base import ModerationAction, ModerationError +from core.moderation.factory import ModerationFactory +from core.ops.entities.trace_entity import TraceTaskName +from core.ops.ops_trace_manager import TraceQueueManager, TraceTask +from core.ops.utils import measure_time + +logger = logging.getLogger(__name__) + + +class InputModeration: + def check( + self, + app_id: str, + tenant_id: str, + app_config: AppConfig, + inputs: Mapping[str, Any], + query: str, + message_id: str, + trace_manager: Optional[TraceQueueManager] = None, + ) -> tuple[bool, Mapping[str, Any], str]: + """ + Process sensitive_word_avoidance. + :param app_id: app id + :param tenant_id: tenant id + :param app_config: app config + :param inputs: inputs + :param query: query + :param message_id: message id + :param trace_manager: trace manager + :return: + """ + inputs = dict(inputs) + if not app_config.sensitive_word_avoidance: + return False, inputs, query + + sensitive_word_avoidance_config = app_config.sensitive_word_avoidance + moderation_type = sensitive_word_avoidance_config.type + + moderation_factory = ModerationFactory( + name=moderation_type, app_id=app_id, tenant_id=tenant_id, config=sensitive_word_avoidance_config.config + ) + + with measure_time() as timer: + moderation_result = moderation_factory.moderation_for_inputs(inputs, query) + + if trace_manager: + trace_manager.add_trace_task( + TraceTask( + TraceTaskName.MODERATION_TRACE, + message_id=message_id, + moderation_result=moderation_result, + inputs=inputs, + timer=timer, + ) + ) + + if not moderation_result.flagged: + return False, inputs, query + + if moderation_result.action == ModerationAction.DIRECT_OUTPUT: + raise ModerationError(moderation_result.preset_response) + elif moderation_result.action == ModerationAction.OVERRIDDEN: + inputs = moderation_result.inputs + query = moderation_result.query + + return True, inputs, query diff --git a/api/core/moderation/keywords/__builtin__ b/api/core/moderation/keywords/__builtin__ new file mode 100644 index 0000000000000000000000000000000000000000..d8263ee9860594d2806b0dfd1bfd17528b0ba2a4 --- /dev/null +++ b/api/core/moderation/keywords/__builtin__ @@ -0,0 +1 @@ +2 \ No newline at end of file diff --git a/api/core/moderation/keywords/__init__.py b/api/core/moderation/keywords/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/moderation/keywords/keywords.py b/api/core/moderation/keywords/keywords.py new file mode 100644 index 0000000000000000000000000000000000000000..9dd2665c3bf3d3f277b3960c8db515ac55b644ed --- /dev/null +++ b/api/core/moderation/keywords/keywords.py @@ -0,0 +1,73 @@ +from collections.abc import Sequence +from typing import Any + +from core.moderation.base import Moderation, ModerationAction, ModerationInputsResult, ModerationOutputsResult + + +class KeywordsModeration(Moderation): + name: str = "keywords" + + @classmethod + def validate_config(cls, tenant_id: str, config: dict) -> None: + """ + Validate the incoming form config data. + + :param tenant_id: the id of workspace + :param config: the form config data + :return: + """ + cls._validate_inputs_and_outputs_config(config, True) + + if not config.get("keywords"): + raise ValueError("keywords is required") + + if len(config.get("keywords", [])) > 10000: + raise ValueError("keywords length must be less than 10000") + + keywords_row_len = config["keywords"].split("\n") + if len(keywords_row_len) > 100: + raise ValueError("the number of rows for the keywords must be less than 100") + + def moderation_for_inputs(self, inputs: dict, query: str = "") -> ModerationInputsResult: + flagged = False + preset_response = "" + if self.config is None: + raise ValueError("The config is not set.") + + if self.config["inputs_config"]["enabled"]: + preset_response = self.config["inputs_config"]["preset_response"] + + if query: + inputs["query__"] = query + + # Filter out empty values + keywords_list = [keyword for keyword in self.config["keywords"].split("\n") if keyword] + + flagged = self._is_violated(inputs, keywords_list) + + return ModerationInputsResult( + flagged=flagged, action=ModerationAction.DIRECT_OUTPUT, preset_response=preset_response + ) + + def moderation_for_outputs(self, text: str) -> ModerationOutputsResult: + flagged = False + preset_response = "" + if self.config is None: + raise ValueError("The config is not set.") + + if self.config["outputs_config"]["enabled"]: + # Filter out empty values + keywords_list = [keyword for keyword in self.config["keywords"].split("\n") if keyword] + + flagged = self._is_violated({"text": text}, keywords_list) + preset_response = self.config["outputs_config"]["preset_response"] + + return ModerationOutputsResult( + flagged=flagged, action=ModerationAction.DIRECT_OUTPUT, preset_response=preset_response + ) + + def _is_violated(self, inputs: dict, keywords_list: list) -> bool: + return any(self._check_keywords_in_value(keywords_list, value) for value in inputs.values()) + + def _check_keywords_in_value(self, keywords_list: Sequence[str], value: Any) -> bool: + return any(keyword.lower() in str(value).lower() for keyword in keywords_list) diff --git a/api/core/moderation/openai_moderation/__builtin__ b/api/core/moderation/openai_moderation/__builtin__ new file mode 100644 index 0000000000000000000000000000000000000000..56a6051ca2b02b04ef92d5150c9ef600403cb1de --- /dev/null +++ b/api/core/moderation/openai_moderation/__builtin__ @@ -0,0 +1 @@ +1 \ No newline at end of file diff --git a/api/core/moderation/openai_moderation/__init__.py b/api/core/moderation/openai_moderation/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/moderation/openai_moderation/openai_moderation.py b/api/core/moderation/openai_moderation/openai_moderation.py new file mode 100644 index 0000000000000000000000000000000000000000..d64f17b383e0b503973f93d9e61151c9adad0b7d --- /dev/null +++ b/api/core/moderation/openai_moderation/openai_moderation.py @@ -0,0 +1,60 @@ +from core.model_manager import ModelManager +from core.model_runtime.entities.model_entities import ModelType +from core.moderation.base import Moderation, ModerationAction, ModerationInputsResult, ModerationOutputsResult + + +class OpenAIModeration(Moderation): + name: str = "openai_moderation" + + @classmethod + def validate_config(cls, tenant_id: str, config: dict) -> None: + """ + Validate the incoming form config data. + + :param tenant_id: the id of workspace + :param config: the form config data + :return: + """ + cls._validate_inputs_and_outputs_config(config, True) + + def moderation_for_inputs(self, inputs: dict, query: str = "") -> ModerationInputsResult: + flagged = False + preset_response = "" + if self.config is None: + raise ValueError("The config is not set.") + + if self.config["inputs_config"]["enabled"]: + preset_response = self.config["inputs_config"]["preset_response"] + + if query: + inputs["query__"] = query + flagged = self._is_violated(inputs) + + return ModerationInputsResult( + flagged=flagged, action=ModerationAction.DIRECT_OUTPUT, preset_response=preset_response + ) + + def moderation_for_outputs(self, text: str) -> ModerationOutputsResult: + flagged = False + preset_response = "" + if self.config is None: + raise ValueError("The config is not set.") + + if self.config["outputs_config"]["enabled"]: + flagged = self._is_violated({"text": text}) + preset_response = self.config["outputs_config"]["preset_response"] + + return ModerationOutputsResult( + flagged=flagged, action=ModerationAction.DIRECT_OUTPUT, preset_response=preset_response + ) + + def _is_violated(self, inputs: dict): + text = "\n".join(str(inputs.values())) + model_manager = ModelManager() + model_instance = model_manager.get_model_instance( + tenant_id=self.tenant_id, provider="openai", model_type=ModelType.MODERATION, model="text-moderation-stable" + ) + + openai_moderation = model_instance.invoke_moderation(text=text) + + return openai_moderation diff --git a/api/core/moderation/output_moderation.py b/api/core/moderation/output_moderation.py new file mode 100644 index 0000000000000000000000000000000000000000..e595be126c7824d21ec9a236946eb9ea9de1c70d --- /dev/null +++ b/api/core/moderation/output_moderation.py @@ -0,0 +1,131 @@ +import logging +import threading +import time +from typing import Any, Optional + +from flask import Flask, current_app +from pydantic import BaseModel, ConfigDict + +from configs import dify_config +from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom +from core.app.entities.queue_entities import QueueMessageReplaceEvent +from core.moderation.base import ModerationAction, ModerationOutputsResult +from core.moderation.factory import ModerationFactory + +logger = logging.getLogger(__name__) + + +class ModerationRule(BaseModel): + type: str + config: dict[str, Any] + + +class OutputModeration(BaseModel): + tenant_id: str + app_id: str + + rule: ModerationRule + queue_manager: AppQueueManager + + thread: Optional[threading.Thread] = None + thread_running: bool = True + buffer: str = "" + is_final_chunk: bool = False + final_output: Optional[str] = None + model_config = ConfigDict(arbitrary_types_allowed=True) + + def should_direct_output(self) -> bool: + return self.final_output is not None + + def get_final_output(self) -> str: + return self.final_output or "" + + def append_new_token(self, token: str) -> None: + self.buffer += token + + if not self.thread: + self.thread = self.start_thread() + + def moderation_completion(self, completion: str, public_event: bool = False) -> str: + self.buffer = completion + self.is_final_chunk = True + + result = self.moderation(tenant_id=self.tenant_id, app_id=self.app_id, moderation_buffer=completion) + + if not result or not result.flagged: + return completion + + if result.action == ModerationAction.DIRECT_OUTPUT: + final_output = result.preset_response + else: + final_output = result.text + + if public_event: + self.queue_manager.publish(QueueMessageReplaceEvent(text=final_output), PublishFrom.TASK_PIPELINE) + + return final_output + + def start_thread(self) -> threading.Thread: + buffer_size = dify_config.MODERATION_BUFFER_SIZE + thread = threading.Thread( + target=self.worker, + kwargs={ + "flask_app": current_app._get_current_object(), # type: ignore + "buffer_size": buffer_size if buffer_size > 0 else dify_config.MODERATION_BUFFER_SIZE, + }, + ) + + thread.start() + + return thread + + def stop_thread(self): + if self.thread and self.thread.is_alive(): + self.thread_running = False + + def worker(self, flask_app: Flask, buffer_size: int): + with flask_app.app_context(): + current_length = 0 + while self.thread_running: + moderation_buffer = self.buffer + buffer_length = len(moderation_buffer) + if not self.is_final_chunk: + chunk_length = buffer_length - current_length + if 0 <= chunk_length < buffer_size: + time.sleep(1) + continue + + current_length = buffer_length + + result = self.moderation( + tenant_id=self.tenant_id, app_id=self.app_id, moderation_buffer=moderation_buffer + ) + + if not result or not result.flagged: + continue + + if result.action == ModerationAction.DIRECT_OUTPUT: + final_output = result.preset_response + self.final_output = final_output + else: + final_output = result.text + self.buffer[len(moderation_buffer) :] + + # trigger replace event + if self.thread_running: + self.queue_manager.publish(QueueMessageReplaceEvent(text=final_output), PublishFrom.TASK_PIPELINE) + + if result.action == ModerationAction.DIRECT_OUTPUT: + break + + def moderation(self, tenant_id: str, app_id: str, moderation_buffer: str) -> Optional[ModerationOutputsResult]: + try: + moderation_factory = ModerationFactory( + name=self.rule.type, app_id=app_id, tenant_id=tenant_id, config=self.rule.config + ) + + result: ModerationOutputsResult = moderation_factory.moderation_for_outputs(moderation_buffer) + return result + except Exception as e: + logger.exception(f"Moderation Output error, app_id: {app_id}") + + return None diff --git a/api/core/ops/__init__.py b/api/core/ops/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/ops/base_trace_instance.py b/api/core/ops/base_trace_instance.py new file mode 100644 index 0000000000000000000000000000000000000000..f7b882fc71d48e13a8a2134a27529f3ebd6bc2b0 --- /dev/null +++ b/api/core/ops/base_trace_instance.py @@ -0,0 +1,26 @@ +from abc import ABC, abstractmethod + +from core.ops.entities.config_entity import BaseTracingConfig +from core.ops.entities.trace_entity import BaseTraceInfo + + +class BaseTraceInstance(ABC): + """ + Base trace instance for ops trace services + """ + + @abstractmethod + def __init__(self, trace_config: BaseTracingConfig): + """ + Abstract initializer for the trace instance. + Distribute trace tasks by matching entities + """ + self.trace_config = trace_config + + @abstractmethod + def trace(self, trace_info: BaseTraceInfo): + """ + Abstract method to trace activities. + Subclasses must implement specific tracing logic for activities. + """ + ... diff --git a/api/core/ops/entities/__init__.py b/api/core/ops/entities/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/ops/entities/config_entity.py b/api/core/ops/entities/config_entity.py new file mode 100644 index 0000000000000000000000000000000000000000..b484242b611d79c02926b308d50e6e925d131759 --- /dev/null +++ b/api/core/ops/entities/config_entity.py @@ -0,0 +1,92 @@ +from enum import Enum + +from pydantic import BaseModel, ValidationInfo, field_validator + + +class TracingProviderEnum(Enum): + LANGFUSE = "langfuse" + LANGSMITH = "langsmith" + OPIK = "opik" + + +class BaseTracingConfig(BaseModel): + """ + Base model class for tracing + """ + + ... + + +class LangfuseConfig(BaseTracingConfig): + """ + Model class for Langfuse tracing config. + """ + + public_key: str + secret_key: str + host: str = "https://api.langfuse.com" + + @field_validator("host") + @classmethod + def set_value(cls, v, info: ValidationInfo): + if v is None or v == "": + v = "https://api.langfuse.com" + if not v.startswith("https://") and not v.startswith("http://"): + raise ValueError("host must start with https:// or http://") + + return v + + +class LangSmithConfig(BaseTracingConfig): + """ + Model class for Langsmith tracing config. + """ + + api_key: str + project: str + endpoint: str = "https://api.smith.langchain.com" + + @field_validator("endpoint") + @classmethod + def set_value(cls, v, info: ValidationInfo): + if v is None or v == "": + v = "https://api.smith.langchain.com" + if not v.startswith("https://"): + raise ValueError("endpoint must start with https://") + + return v + + +class OpikConfig(BaseTracingConfig): + """ + Model class for Opik tracing config. + """ + + api_key: str | None = None + project: str | None = None + workspace: str | None = None + url: str = "https://www.comet.com/opik/api/" + + @field_validator("project") + @classmethod + def project_validator(cls, v, info: ValidationInfo): + if v is None or v == "": + v = "Default Project" + + return v + + @field_validator("url") + @classmethod + def url_validator(cls, v, info: ValidationInfo): + if v is None or v == "": + v = "https://www.comet.com/opik/api/" + if not v.startswith(("https://", "http://")): + raise ValueError("url must start with https:// or http://") + if not v.endswith("/api/"): + raise ValueError("url should ends with /api/") + + return v + + +OPS_FILE_PATH = "ops_trace/" +OPS_TRACE_FAILED_KEY = "FAILED_OPS_TRACE" diff --git a/api/core/ops/entities/trace_entity.py b/api/core/ops/entities/trace_entity.py new file mode 100644 index 0000000000000000000000000000000000000000..f0e34c0cd71241508c57c23b9518b1b9255f4a52 --- /dev/null +++ b/api/core/ops/entities/trace_entity.py @@ -0,0 +1,134 @@ +from collections.abc import Mapping +from datetime import datetime +from enum import StrEnum +from typing import Any, Optional, Union + +from pydantic import BaseModel, ConfigDict, field_validator + + +class BaseTraceInfo(BaseModel): + message_id: Optional[str] = None + message_data: Optional[Any] = None + inputs: Optional[Union[str, dict[str, Any], list]] = None + outputs: Optional[Union[str, dict[str, Any], list]] = None + start_time: Optional[datetime] = None + end_time: Optional[datetime] = None + metadata: dict[str, Any] + + @field_validator("inputs", "outputs") + @classmethod + def ensure_type(cls, v): + if v is None: + return None + if isinstance(v, str | dict | list): + return v + return "" + + class Config: + json_encoders = { + datetime: lambda v: v.isoformat(), + } + + +class WorkflowTraceInfo(BaseTraceInfo): + workflow_data: Any + conversation_id: Optional[str] = None + workflow_app_log_id: Optional[str] = None + workflow_id: str + tenant_id: str + workflow_run_id: str + workflow_run_elapsed_time: Union[int, float] + workflow_run_status: str + workflow_run_inputs: Mapping[str, Any] + workflow_run_outputs: Mapping[str, Any] + workflow_run_version: str + error: Optional[str] = None + total_tokens: int + file_list: list[str] + query: str + metadata: dict[str, Any] + + +class MessageTraceInfo(BaseTraceInfo): + conversation_model: str + message_tokens: int + answer_tokens: int + total_tokens: int + error: Optional[str] = None + file_list: Optional[Union[str, dict[str, Any], list]] = None + message_file_data: Optional[Any] = None + conversation_mode: str + + +class ModerationTraceInfo(BaseTraceInfo): + flagged: bool + action: str + preset_response: str + query: str + + +class SuggestedQuestionTraceInfo(BaseTraceInfo): + total_tokens: int + status: Optional[str] = None + error: Optional[str] = None + from_account_id: Optional[str] = None + agent_based: Optional[bool] = None + from_source: Optional[str] = None + model_provider: Optional[str] = None + model_id: Optional[str] = None + suggested_question: list[str] + level: str + status_message: Optional[str] = None + workflow_run_id: Optional[str] = None + + model_config = ConfigDict(protected_namespaces=()) + + +class DatasetRetrievalTraceInfo(BaseTraceInfo): + documents: Any + + +class ToolTraceInfo(BaseTraceInfo): + tool_name: str + tool_inputs: dict[str, Any] + tool_outputs: str + metadata: dict[str, Any] + message_file_data: Any + error: Optional[str] = None + tool_config: dict[str, Any] + time_cost: Union[int, float] + tool_parameters: dict[str, Any] + file_url: Union[str, None, list] + + +class GenerateNameTraceInfo(BaseTraceInfo): + conversation_id: Optional[str] = None + tenant_id: str + + +class TaskData(BaseModel): + app_id: str + trace_info_type: str + trace_info: Any + + +trace_info_info_map = { + "WorkflowTraceInfo": WorkflowTraceInfo, + "MessageTraceInfo": MessageTraceInfo, + "ModerationTraceInfo": ModerationTraceInfo, + "SuggestedQuestionTraceInfo": SuggestedQuestionTraceInfo, + "DatasetRetrievalTraceInfo": DatasetRetrievalTraceInfo, + "ToolTraceInfo": ToolTraceInfo, + "GenerateNameTraceInfo": GenerateNameTraceInfo, +} + + +class TraceTaskName(StrEnum): + CONVERSATION_TRACE = "conversation" + WORKFLOW_TRACE = "workflow" + MESSAGE_TRACE = "message" + MODERATION_TRACE = "moderation" + SUGGESTED_QUESTION_TRACE = "suggested_question" + DATASET_RETRIEVAL_TRACE = "dataset_retrieval" + TOOL_TRACE = "tool" + GENERATE_NAME_TRACE = "generate_conversation_name" diff --git a/api/core/ops/langfuse_trace/__init__.py b/api/core/ops/langfuse_trace/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/ops/langfuse_trace/entities/__init__.py b/api/core/ops/langfuse_trace/entities/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/ops/langfuse_trace/entities/langfuse_trace_entity.py b/api/core/ops/langfuse_trace/entities/langfuse_trace_entity.py new file mode 100644 index 0000000000000000000000000000000000000000..f486da3a6d0c90dfe7c61167fe5b9054e01dbbf9 --- /dev/null +++ b/api/core/ops/langfuse_trace/entities/langfuse_trace_entity.py @@ -0,0 +1,282 @@ +from datetime import datetime +from enum import StrEnum +from typing import Any, Optional, Union + +from pydantic import BaseModel, ConfigDict, Field, field_validator +from pydantic_core.core_schema import ValidationInfo + +from core.ops.utils import replace_text_with_content + + +def validate_input_output(v, field_name): + """ + Validate input output + :param v: + :param field_name: + :return: + """ + if v == {} or v is None: + return v + if isinstance(v, str): + return [ + { + "role": "assistant" if field_name == "output" else "user", + "content": v, + } + ] + elif isinstance(v, list): + if len(v) > 0 and isinstance(v[0], dict): + v = replace_text_with_content(data=v) + return v + else: + return [ + { + "role": "assistant" if field_name == "output" else "user", + "content": str(v), + } + ] + + return v + + +class LevelEnum(StrEnum): + DEBUG = "DEBUG" + WARNING = "WARNING" + ERROR = "ERROR" + DEFAULT = "DEFAULT" + + +class LangfuseTrace(BaseModel): + """ + Langfuse trace model + """ + + id: Optional[str] = Field( + default=None, + description="The id of the trace can be set, defaults to a random id. Used to link traces to external systems " + "or when creating a distributed trace. Traces are upserted on id.", + ) + name: Optional[str] = Field( + default=None, + description="Identifier of the trace. Useful for sorting/filtering in the UI.", + ) + input: Optional[Union[str, dict[str, Any], list, None]] = Field( + default=None, description="The input of the trace. Can be any JSON object." + ) + output: Optional[Union[str, dict[str, Any], list, None]] = Field( + default=None, description="The output of the trace. Can be any JSON object." + ) + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="Additional metadata of the trace. Can be any JSON object. Metadata is merged when being updated " + "via the API.", + ) + user_id: Optional[str] = Field( + default=None, + description="The id of the user that triggered the execution. Used to provide user-level analytics.", + ) + session_id: Optional[str] = Field( + default=None, + description="Used to group multiple traces into a session in Langfuse. Use your own session/thread identifier.", + ) + version: Optional[str] = Field( + default=None, + description="The version of the trace type. Used to understand how changes to the trace type affect metrics. " + "Useful in debugging.", + ) + release: Optional[str] = Field( + default=None, + description="The release identifier of the current deployment. Used to understand how changes of different " + "deployments affect metrics. Useful in debugging.", + ) + tags: Optional[list[str]] = Field( + default=None, + description="Tags are used to categorize or label traces. Traces can be filtered by tags in the UI and GET " + "API. Tags can also be changed in the UI. Tags are merged and never deleted via the API.", + ) + public: Optional[bool] = Field( + default=None, + description="You can make a trace public to share it via a public link. This allows others to view the trace " + "without needing to log in or be members of your Langfuse project.", + ) + + @field_validator("input", "output") + @classmethod + def ensure_dict(cls, v, info: ValidationInfo): + field_name = info.field_name + return validate_input_output(v, field_name) + + +class LangfuseSpan(BaseModel): + """ + Langfuse span model + """ + + id: Optional[str] = Field( + default=None, + description="The id of the span can be set, otherwise a random id is generated. Spans are upserted on id.", + ) + session_id: Optional[str] = Field( + default=None, + description="Used to group multiple spans into a session in Langfuse. Use your own session/thread identifier.", + ) + trace_id: Optional[str] = Field( + default=None, + description="The id of the trace the span belongs to. Used to link spans to traces.", + ) + user_id: Optional[str] = Field( + default=None, + description="The id of the user that triggered the execution. Used to provide user-level analytics.", + ) + start_time: Optional[datetime | str] = Field( + default_factory=datetime.now, + description="The time at which the span started, defaults to the current time.", + ) + end_time: Optional[datetime | str] = Field( + default=None, + description="The time at which the span ended. Automatically set by span.end().", + ) + name: Optional[str] = Field( + default=None, + description="Identifier of the span. Useful for sorting/filtering in the UI.", + ) + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="Additional metadata of the span. Can be any JSON object. Metadata is merged when being updated " + "via the API.", + ) + level: Optional[str] = Field( + default=None, + description="The level of the span. Can be DEBUG, DEFAULT, WARNING or ERROR. Used for sorting/filtering of " + "traces with elevated error levels and for highlighting in the UI.", + ) + status_message: Optional[str] = Field( + default=None, + description="The status message of the span. Additional field for context of the event. E.g. the error " + "message of an error event.", + ) + input: Optional[Union[str, dict[str, Any], list, None]] = Field( + default=None, description="The input of the span. Can be any JSON object." + ) + output: Optional[Union[str, dict[str, Any], list, None]] = Field( + default=None, description="The output of the span. Can be any JSON object." + ) + version: Optional[str] = Field( + default=None, + description="The version of the span type. Used to understand how changes to the span type affect metrics. " + "Useful in debugging.", + ) + parent_observation_id: Optional[str] = Field( + default=None, + description="The id of the observation the span belongs to. Used to link spans to observations.", + ) + + @field_validator("input", "output") + @classmethod + def ensure_dict(cls, v, info: ValidationInfo): + field_name = info.field_name + return validate_input_output(v, field_name) + + +class UnitEnum(StrEnum): + CHARACTERS = "CHARACTERS" + TOKENS = "TOKENS" + SECONDS = "SECONDS" + MILLISECONDS = "MILLISECONDS" + IMAGES = "IMAGES" + + +class GenerationUsage(BaseModel): + promptTokens: Optional[int] = None + completionTokens: Optional[int] = None + total: Optional[int] = None + input: Optional[int] = None + output: Optional[int] = None + unit: Optional[UnitEnum] = None + inputCost: Optional[float] = None + outputCost: Optional[float] = None + totalCost: Optional[float] = None + + @field_validator("input", "output") + @classmethod + def ensure_dict(cls, v, info: ValidationInfo): + field_name = info.field_name + return validate_input_output(v, field_name) + + +class LangfuseGeneration(BaseModel): + id: Optional[str] = Field( + default=None, + description="The id of the generation can be set, defaults to random id.", + ) + trace_id: Optional[str] = Field( + default=None, + description="The id of the trace the generation belongs to. Used to link generations to traces.", + ) + parent_observation_id: Optional[str] = Field( + default=None, + description="The id of the observation the generation belongs to. Used to link generations to observations.", + ) + name: Optional[str] = Field( + default=None, + description="Identifier of the generation. Useful for sorting/filtering in the UI.", + ) + start_time: Optional[datetime | str] = Field( + default_factory=datetime.now, + description="The time at which the generation started, defaults to the current time.", + ) + completion_start_time: Optional[datetime | str] = Field( + default=None, + description="The time at which the completion started (streaming). Set it to get latency analytics broken " + "down into time until completion started and completion duration.", + ) + end_time: Optional[datetime | str] = Field( + default=None, + description="The time at which the generation ended. Automatically set by generation.end().", + ) + model: Optional[str] = Field(default=None, description="The name of the model used for the generation.") + model_parameters: Optional[dict[str, Any]] = Field( + default=None, + description="The parameters of the model used for the generation; can be any key-value pairs.", + ) + input: Optional[Any] = Field( + default=None, + description="The prompt used for the generation. Can be any string or JSON object.", + ) + output: Optional[Any] = Field( + default=None, + description="The completion generated by the model. Can be any string or JSON object.", + ) + usage: Optional[GenerationUsage] = Field( + default=None, + description="The usage object supports the OpenAi structure with tokens and a more generic version with " + "detailed costs and units.", + ) + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="Additional metadata of the generation. Can be any JSON object. Metadata is merged when being " + "updated via the API.", + ) + level: Optional[LevelEnum] = Field( + default=None, + description="The level of the generation. Can be DEBUG, DEFAULT, WARNING or ERROR. Used for sorting/filtering " + "of traces with elevated error levels and for highlighting in the UI.", + ) + status_message: Optional[str] = Field( + default=None, + description="The status message of the generation. Additional field for context of the event. E.g. the error " + "message of an error event.", + ) + version: Optional[str] = Field( + default=None, + description="The version of the generation type. Used to understand how changes to the span type affect " + "metrics. Useful in debugging.", + ) + + model_config = ConfigDict(protected_namespaces=()) + + @field_validator("input", "output") + @classmethod + def ensure_dict(cls, v, info: ValidationInfo): + field_name = info.field_name + return validate_input_output(v, field_name) diff --git a/api/core/ops/langfuse_trace/langfuse_trace.py b/api/core/ops/langfuse_trace/langfuse_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..b9ba068b19936d0c8e27174610dc14e158020a70 --- /dev/null +++ b/api/core/ops/langfuse_trace/langfuse_trace.py @@ -0,0 +1,455 @@ +import json +import logging +import os +from datetime import datetime, timedelta +from typing import Optional + +from langfuse import Langfuse # type: ignore + +from core.ops.base_trace_instance import BaseTraceInstance +from core.ops.entities.config_entity import LangfuseConfig +from core.ops.entities.trace_entity import ( + BaseTraceInfo, + DatasetRetrievalTraceInfo, + GenerateNameTraceInfo, + MessageTraceInfo, + ModerationTraceInfo, + SuggestedQuestionTraceInfo, + ToolTraceInfo, + TraceTaskName, + WorkflowTraceInfo, +) +from core.ops.langfuse_trace.entities.langfuse_trace_entity import ( + GenerationUsage, + LangfuseGeneration, + LangfuseSpan, + LangfuseTrace, + LevelEnum, + UnitEnum, +) +from core.ops.utils import filter_none_values +from extensions.ext_database import db +from models.model import EndUser +from models.workflow import WorkflowNodeExecution + +logger = logging.getLogger(__name__) + + +class LangFuseDataTrace(BaseTraceInstance): + def __init__( + self, + langfuse_config: LangfuseConfig, + ): + super().__init__(langfuse_config) + self.langfuse_client = Langfuse( + public_key=langfuse_config.public_key, + secret_key=langfuse_config.secret_key, + host=langfuse_config.host, + ) + self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001") + + def trace(self, trace_info: BaseTraceInfo): + if isinstance(trace_info, WorkflowTraceInfo): + self.workflow_trace(trace_info) + if isinstance(trace_info, MessageTraceInfo): + self.message_trace(trace_info) + if isinstance(trace_info, ModerationTraceInfo): + self.moderation_trace(trace_info) + if isinstance(trace_info, SuggestedQuestionTraceInfo): + self.suggested_question_trace(trace_info) + if isinstance(trace_info, DatasetRetrievalTraceInfo): + self.dataset_retrieval_trace(trace_info) + if isinstance(trace_info, ToolTraceInfo): + self.tool_trace(trace_info) + if isinstance(trace_info, GenerateNameTraceInfo): + self.generate_name_trace(trace_info) + + def workflow_trace(self, trace_info: WorkflowTraceInfo): + trace_id = trace_info.workflow_run_id + user_id = trace_info.metadata.get("user_id") + metadata = trace_info.metadata + metadata["workflow_app_log_id"] = trace_info.workflow_app_log_id + + if trace_info.message_id: + trace_id = trace_info.message_id + name = TraceTaskName.MESSAGE_TRACE.value + trace_data = LangfuseTrace( + id=trace_id, + user_id=user_id, + name=name, + input=dict(trace_info.workflow_run_inputs), + output=dict(trace_info.workflow_run_outputs), + metadata=metadata, + session_id=trace_info.conversation_id, + tags=["message", "workflow"], + ) + self.add_trace(langfuse_trace_data=trace_data) + workflow_span_data = LangfuseSpan( + id=trace_info.workflow_run_id, + name=TraceTaskName.WORKFLOW_TRACE.value, + input=dict(trace_info.workflow_run_inputs), + output=dict(trace_info.workflow_run_outputs), + trace_id=trace_id, + start_time=trace_info.start_time, + end_time=trace_info.end_time, + metadata=metadata, + level=LevelEnum.DEFAULT if trace_info.error == "" else LevelEnum.ERROR, + status_message=trace_info.error or "", + ) + self.add_span(langfuse_span_data=workflow_span_data) + else: + trace_data = LangfuseTrace( + id=trace_id, + user_id=user_id, + name=TraceTaskName.WORKFLOW_TRACE.value, + input=dict(trace_info.workflow_run_inputs), + output=dict(trace_info.workflow_run_outputs), + metadata=metadata, + session_id=trace_info.conversation_id, + tags=["workflow"], + ) + self.add_trace(langfuse_trace_data=trace_data) + + # through workflow_run_id get all_nodes_execution + workflow_nodes_execution_id_records = ( + db.session.query(WorkflowNodeExecution.id) + .filter(WorkflowNodeExecution.workflow_run_id == trace_info.workflow_run_id) + .all() + ) + + for node_execution_id_record in workflow_nodes_execution_id_records: + node_execution = ( + db.session.query( + WorkflowNodeExecution.id, + WorkflowNodeExecution.tenant_id, + WorkflowNodeExecution.app_id, + WorkflowNodeExecution.title, + WorkflowNodeExecution.node_type, + WorkflowNodeExecution.status, + WorkflowNodeExecution.inputs, + WorkflowNodeExecution.outputs, + WorkflowNodeExecution.created_at, + WorkflowNodeExecution.elapsed_time, + WorkflowNodeExecution.process_data, + WorkflowNodeExecution.execution_metadata, + ) + .filter(WorkflowNodeExecution.id == node_execution_id_record.id) + .first() + ) + + if not node_execution: + continue + + node_execution_id = node_execution.id + tenant_id = node_execution.tenant_id + app_id = node_execution.app_id + node_name = node_execution.title + node_type = node_execution.node_type + status = node_execution.status + if node_type == "llm": + inputs = ( + json.loads(node_execution.process_data).get("prompts", {}) if node_execution.process_data else {} + ) + else: + inputs = json.loads(node_execution.inputs) if node_execution.inputs else {} + outputs = json.loads(node_execution.outputs) if node_execution.outputs else {} + created_at = node_execution.created_at or datetime.now() + elapsed_time = node_execution.elapsed_time + finished_at = created_at + timedelta(seconds=elapsed_time) + + metadata = json.loads(node_execution.execution_metadata) if node_execution.execution_metadata else {} + metadata.update( + { + "workflow_run_id": trace_info.workflow_run_id, + "node_execution_id": node_execution_id, + "tenant_id": tenant_id, + "app_id": app_id, + "node_name": node_name, + "node_type": node_type, + "status": status, + } + ) + process_data = json.loads(node_execution.process_data) if node_execution.process_data else {} + model_provider = process_data.get("model_provider", None) + model_name = process_data.get("model_name", None) + if model_provider is not None and model_name is not None: + metadata.update( + { + "model_provider": model_provider, + "model_name": model_name, + } + ) + + # add span + if trace_info.message_id: + span_data = LangfuseSpan( + id=node_execution_id, + name=node_type, + input=inputs, + output=outputs, + trace_id=trace_id, + start_time=created_at, + end_time=finished_at, + metadata=metadata, + level=(LevelEnum.DEFAULT if status == "succeeded" else LevelEnum.ERROR), + status_message=trace_info.error or "", + parent_observation_id=trace_info.workflow_run_id, + ) + else: + span_data = LangfuseSpan( + id=node_execution_id, + name=node_type, + input=inputs, + output=outputs, + trace_id=trace_id, + start_time=created_at, + end_time=finished_at, + metadata=metadata, + level=(LevelEnum.DEFAULT if status == "succeeded" else LevelEnum.ERROR), + status_message=trace_info.error or "", + ) + + self.add_span(langfuse_span_data=span_data) + + if process_data and process_data.get("model_mode") == "chat": + total_token = metadata.get("total_tokens", 0) + # add generation + generation_usage = GenerationUsage( + total=total_token, + ) + + node_generation_data = LangfuseGeneration( + name="llm", + trace_id=trace_id, + model=process_data.get("model_name"), + parent_observation_id=node_execution_id, + start_time=created_at, + end_time=finished_at, + input=inputs, + output=outputs, + metadata=metadata, + level=(LevelEnum.DEFAULT if status == "succeeded" else LevelEnum.ERROR), + status_message=trace_info.error or "", + usage=generation_usage, + ) + + self.add_generation(langfuse_generation_data=node_generation_data) + + def message_trace(self, trace_info: MessageTraceInfo, **kwargs): + # get message file data + file_list = trace_info.file_list + metadata = trace_info.metadata + message_data = trace_info.message_data + if message_data is None: + return + message_id = message_data.id + + user_id = message_data.from_account_id + if message_data.from_end_user_id: + end_user_data: Optional[EndUser] = ( + db.session.query(EndUser).filter(EndUser.id == message_data.from_end_user_id).first() + ) + if end_user_data is not None: + user_id = end_user_data.session_id + metadata["user_id"] = user_id + + trace_data = LangfuseTrace( + id=message_id, + user_id=user_id, + name=TraceTaskName.MESSAGE_TRACE.value, + input={ + "message": trace_info.inputs, + "files": file_list, + "message_tokens": trace_info.message_tokens, + "answer_tokens": trace_info.answer_tokens, + "total_tokens": trace_info.total_tokens, + "error": trace_info.error, + "provider_response_latency": message_data.provider_response_latency, + "created_at": trace_info.start_time, + }, + output=trace_info.outputs, + metadata=metadata, + session_id=message_data.conversation_id, + tags=["message", str(trace_info.conversation_mode)], + version=None, + release=None, + public=None, + ) + self.add_trace(langfuse_trace_data=trace_data) + + # start add span + generation_usage = GenerationUsage( + input=trace_info.message_tokens, + output=trace_info.answer_tokens, + total=trace_info.total_tokens, + unit=UnitEnum.TOKENS, + totalCost=message_data.total_price, + ) + + langfuse_generation_data = LangfuseGeneration( + name="llm", + trace_id=message_id, + start_time=trace_info.start_time, + end_time=trace_info.end_time, + model=message_data.model_id, + input=trace_info.inputs, + output=message_data.answer, + metadata=metadata, + level=(LevelEnum.DEFAULT if message_data.status != "error" else LevelEnum.ERROR), + status_message=message_data.error or "", + usage=generation_usage, + ) + + self.add_generation(langfuse_generation_data) + + def moderation_trace(self, trace_info: ModerationTraceInfo): + if trace_info.message_data is None: + return + span_data = LangfuseSpan( + name=TraceTaskName.MODERATION_TRACE.value, + input=trace_info.inputs, + output={ + "action": trace_info.action, + "flagged": trace_info.flagged, + "preset_response": trace_info.preset_response, + "inputs": trace_info.inputs, + }, + trace_id=trace_info.message_id, + start_time=trace_info.start_time or trace_info.message_data.created_at, + end_time=trace_info.end_time or trace_info.message_data.created_at, + metadata=trace_info.metadata, + ) + + self.add_span(langfuse_span_data=span_data) + + def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo): + message_data = trace_info.message_data + if message_data is None: + return + generation_usage = GenerationUsage( + total=len(str(trace_info.suggested_question)), + input=len(trace_info.inputs) if trace_info.inputs else 0, + output=len(trace_info.suggested_question), + unit=UnitEnum.CHARACTERS, + ) + + generation_data = LangfuseGeneration( + name=TraceTaskName.SUGGESTED_QUESTION_TRACE.value, + input=trace_info.inputs, + output=str(trace_info.suggested_question), + trace_id=trace_info.message_id, + start_time=trace_info.start_time, + end_time=trace_info.end_time, + metadata=trace_info.metadata, + level=(LevelEnum.DEFAULT if message_data.status != "error" else LevelEnum.ERROR), + status_message=message_data.error or "", + usage=generation_usage, + ) + + self.add_generation(langfuse_generation_data=generation_data) + + def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo): + if trace_info.message_data is None: + return + dataset_retrieval_span_data = LangfuseSpan( + name=TraceTaskName.DATASET_RETRIEVAL_TRACE.value, + input=trace_info.inputs, + output={"documents": trace_info.documents}, + trace_id=trace_info.message_id, + start_time=trace_info.start_time or trace_info.message_data.created_at, + end_time=trace_info.end_time or trace_info.message_data.updated_at, + metadata=trace_info.metadata, + ) + + self.add_span(langfuse_span_data=dataset_retrieval_span_data) + + def tool_trace(self, trace_info: ToolTraceInfo): + tool_span_data = LangfuseSpan( + name=trace_info.tool_name, + input=trace_info.tool_inputs, + output=trace_info.tool_outputs, + trace_id=trace_info.message_id, + start_time=trace_info.start_time, + end_time=trace_info.end_time, + metadata=trace_info.metadata, + level=(LevelEnum.DEFAULT if trace_info.error == "" or trace_info.error is None else LevelEnum.ERROR), + status_message=trace_info.error, + ) + + self.add_span(langfuse_span_data=tool_span_data) + + def generate_name_trace(self, trace_info: GenerateNameTraceInfo): + name_generation_trace_data = LangfuseTrace( + name=TraceTaskName.GENERATE_NAME_TRACE.value, + input=trace_info.inputs, + output=trace_info.outputs, + user_id=trace_info.tenant_id, + metadata=trace_info.metadata, + session_id=trace_info.conversation_id, + ) + + self.add_trace(langfuse_trace_data=name_generation_trace_data) + + name_generation_span_data = LangfuseSpan( + name=TraceTaskName.GENERATE_NAME_TRACE.value, + input=trace_info.inputs, + output=trace_info.outputs, + trace_id=trace_info.conversation_id, + start_time=trace_info.start_time, + end_time=trace_info.end_time, + metadata=trace_info.metadata, + ) + self.add_span(langfuse_span_data=name_generation_span_data) + + def add_trace(self, langfuse_trace_data: Optional[LangfuseTrace] = None): + format_trace_data = filter_none_values(langfuse_trace_data.model_dump()) if langfuse_trace_data else {} + try: + self.langfuse_client.trace(**format_trace_data) + logger.debug("LangFuse Trace created successfully") + except Exception as e: + raise ValueError(f"LangFuse Failed to create trace: {str(e)}") + + def add_span(self, langfuse_span_data: Optional[LangfuseSpan] = None): + format_span_data = filter_none_values(langfuse_span_data.model_dump()) if langfuse_span_data else {} + try: + self.langfuse_client.span(**format_span_data) + logger.debug("LangFuse Span created successfully") + except Exception as e: + raise ValueError(f"LangFuse Failed to create span: {str(e)}") + + def update_span(self, span, langfuse_span_data: Optional[LangfuseSpan] = None): + format_span_data = filter_none_values(langfuse_span_data.model_dump()) if langfuse_span_data else {} + + span.end(**format_span_data) + + def add_generation(self, langfuse_generation_data: Optional[LangfuseGeneration] = None): + format_generation_data = ( + filter_none_values(langfuse_generation_data.model_dump()) if langfuse_generation_data else {} + ) + try: + self.langfuse_client.generation(**format_generation_data) + logger.debug("LangFuse Generation created successfully") + except Exception as e: + raise ValueError(f"LangFuse Failed to create generation: {str(e)}") + + def update_generation(self, generation, langfuse_generation_data: Optional[LangfuseGeneration] = None): + format_generation_data = ( + filter_none_values(langfuse_generation_data.model_dump()) if langfuse_generation_data else {} + ) + + generation.end(**format_generation_data) + + def api_check(self): + try: + return self.langfuse_client.auth_check() + except Exception as e: + logger.debug(f"LangFuse API check failed: {str(e)}") + raise ValueError(f"LangFuse API check failed: {str(e)}") + + def get_project_key(self): + try: + projects = self.langfuse_client.client.projects.get() + return projects.data[0].id + except Exception as e: + logger.debug(f"LangFuse get project key failed: {str(e)}") + raise ValueError(f"LangFuse get project key failed: {str(e)}") diff --git a/api/core/ops/langsmith_trace/__init__.py b/api/core/ops/langsmith_trace/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/ops/langsmith_trace/entities/__init__.py b/api/core/ops/langsmith_trace/entities/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/ops/langsmith_trace/entities/langsmith_trace_entity.py b/api/core/ops/langsmith_trace/entities/langsmith_trace_entity.py new file mode 100644 index 0000000000000000000000000000000000000000..348b7ba5012b6bca84fe30c34e9322bece4a6b54 --- /dev/null +++ b/api/core/ops/langsmith_trace/entities/langsmith_trace_entity.py @@ -0,0 +1,141 @@ +from datetime import datetime +from enum import StrEnum +from typing import Any, Optional, Union + +from pydantic import BaseModel, Field, field_validator +from pydantic_core.core_schema import ValidationInfo + +from core.ops.utils import replace_text_with_content + + +class LangSmithRunType(StrEnum): + tool = "tool" + chain = "chain" + llm = "llm" + retriever = "retriever" + embedding = "embedding" + prompt = "prompt" + parser = "parser" + + +class LangSmithTokenUsage(BaseModel): + input_tokens: Optional[int] = None + output_tokens: Optional[int] = None + total_tokens: Optional[int] = None + + +class LangSmithMultiModel(BaseModel): + file_list: Optional[list[str]] = Field(None, description="List of files") + + +class LangSmithRunModel(LangSmithTokenUsage, LangSmithMultiModel): + name: Optional[str] = Field(..., description="Name of the run") + inputs: Optional[Union[str, dict[str, Any], list, None]] = Field(None, description="Inputs of the run") + outputs: Optional[Union[str, dict[str, Any], list, None]] = Field(None, description="Outputs of the run") + run_type: LangSmithRunType = Field(..., description="Type of the run") + start_time: Optional[datetime | str] = Field(None, description="Start time of the run") + end_time: Optional[datetime | str] = Field(None, description="End time of the run") + extra: Optional[dict[str, Any]] = Field(None, description="Extra information of the run") + error: Optional[str] = Field(None, description="Error message of the run") + serialized: Optional[dict[str, Any]] = Field(None, description="Serialized data of the run") + parent_run_id: Optional[str] = Field(None, description="Parent run ID") + events: Optional[list[dict[str, Any]]] = Field(None, description="Events associated with the run") + tags: Optional[list[str]] = Field(None, description="Tags associated with the run") + trace_id: Optional[str] = Field(None, description="Trace ID associated with the run") + dotted_order: Optional[str] = Field(None, description="Dotted order of the run") + id: Optional[str] = Field(None, description="ID of the run") + session_id: Optional[str] = Field(None, description="Session ID associated with the run") + session_name: Optional[str] = Field(None, description="Session name associated with the run") + reference_example_id: Optional[str] = Field(None, description="Reference example ID associated with the run") + input_attachments: Optional[dict[str, Any]] = Field(None, description="Input attachments of the run") + output_attachments: Optional[dict[str, Any]] = Field(None, description="Output attachments of the run") + + @field_validator("inputs", "outputs") + @classmethod + def ensure_dict(cls, v, info: ValidationInfo): + field_name = info.field_name + values = info.data + if v == {} or v is None: + return v + usage_metadata = { + "input_tokens": values.get("input_tokens", 0), + "output_tokens": values.get("output_tokens", 0), + "total_tokens": values.get("total_tokens", 0), + } + file_list = values.get("file_list", []) + if isinstance(v, str): + if field_name == "inputs": + return { + "messages": { + "role": "user", + "content": v, + "usage_metadata": usage_metadata, + "file_list": file_list, + }, + } + elif field_name == "outputs": + return { + "choices": { + "role": "ai", + "content": v, + "usage_metadata": usage_metadata, + "file_list": file_list, + }, + } + elif isinstance(v, list): + data = {} + if len(v) > 0 and isinstance(v[0], dict): + # rename text to content + v = replace_text_with_content(data=v) + if field_name == "inputs": + data = { + "messages": v, + } + elif field_name == "outputs": + data = { + "choices": { + "role": "ai", + "content": v, + "usage_metadata": usage_metadata, + "file_list": file_list, + }, + } + return data + else: + return { + "choices": { + "role": "ai" if field_name == "outputs" else "user", + "content": str(v), + "usage_metadata": usage_metadata, + "file_list": file_list, + }, + } + if isinstance(v, dict): + v["usage_metadata"] = usage_metadata + v["file_list"] = file_list + return v + return v + + @classmethod + @field_validator("start_time", "end_time") + def format_time(cls, v, info: ValidationInfo): + if not isinstance(v, datetime): + raise ValueError(f"{info.field_name} must be a datetime object") + else: + return v.strftime("%Y-%m-%dT%H:%M:%S.%fZ") + + +class LangSmithRunUpdateModel(BaseModel): + run_id: str = Field(..., description="ID of the run") + trace_id: Optional[str] = Field(None, description="Trace ID associated with the run") + dotted_order: Optional[str] = Field(None, description="Dotted order of the run") + parent_run_id: Optional[str] = Field(None, description="Parent run ID") + end_time: Optional[datetime | str] = Field(None, description="End time of the run") + error: Optional[str] = Field(None, description="Error message of the run") + inputs: Optional[dict[str, Any]] = Field(None, description="Inputs of the run") + outputs: Optional[dict[str, Any]] = Field(None, description="Outputs of the run") + events: Optional[list[dict[str, Any]]] = Field(None, description="Events associated with the run") + tags: Optional[list[str]] = Field(None, description="Tags associated with the run") + extra: Optional[dict[str, Any]] = Field(None, description="Extra information of the run") + input_attachments: Optional[dict[str, Any]] = Field(None, description="Input attachments of the run") + output_attachments: Optional[dict[str, Any]] = Field(None, description="Output attachments of the run") diff --git a/api/core/ops/langsmith_trace/langsmith_trace.py b/api/core/ops/langsmith_trace/langsmith_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..4ffd888bddf8a39f2f282b22c72fb00caaad0299 --- /dev/null +++ b/api/core/ops/langsmith_trace/langsmith_trace.py @@ -0,0 +1,524 @@ +import json +import logging +import os +import uuid +from datetime import datetime, timedelta +from typing import Optional, cast + +from langsmith import Client +from langsmith.schemas import RunBase + +from core.ops.base_trace_instance import BaseTraceInstance +from core.ops.entities.config_entity import LangSmithConfig +from core.ops.entities.trace_entity import ( + BaseTraceInfo, + DatasetRetrievalTraceInfo, + GenerateNameTraceInfo, + MessageTraceInfo, + ModerationTraceInfo, + SuggestedQuestionTraceInfo, + ToolTraceInfo, + TraceTaskName, + WorkflowTraceInfo, +) +from core.ops.langsmith_trace.entities.langsmith_trace_entity import ( + LangSmithRunModel, + LangSmithRunType, + LangSmithRunUpdateModel, +) +from core.ops.utils import filter_none_values, generate_dotted_order +from extensions.ext_database import db +from models.model import EndUser, MessageFile +from models.workflow import WorkflowNodeExecution + +logger = logging.getLogger(__name__) + + +class LangSmithDataTrace(BaseTraceInstance): + def __init__( + self, + langsmith_config: LangSmithConfig, + ): + super().__init__(langsmith_config) + self.langsmith_key = langsmith_config.api_key + self.project_name = langsmith_config.project + self.project_id = None + self.langsmith_client = Client(api_key=langsmith_config.api_key, api_url=langsmith_config.endpoint) + self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001") + + def trace(self, trace_info: BaseTraceInfo): + if isinstance(trace_info, WorkflowTraceInfo): + self.workflow_trace(trace_info) + if isinstance(trace_info, MessageTraceInfo): + self.message_trace(trace_info) + if isinstance(trace_info, ModerationTraceInfo): + self.moderation_trace(trace_info) + if isinstance(trace_info, SuggestedQuestionTraceInfo): + self.suggested_question_trace(trace_info) + if isinstance(trace_info, DatasetRetrievalTraceInfo): + self.dataset_retrieval_trace(trace_info) + if isinstance(trace_info, ToolTraceInfo): + self.tool_trace(trace_info) + if isinstance(trace_info, GenerateNameTraceInfo): + self.generate_name_trace(trace_info) + + def workflow_trace(self, trace_info: WorkflowTraceInfo): + trace_id = trace_info.message_id or trace_info.workflow_run_id + if trace_info.start_time is None: + trace_info.start_time = datetime.now() + message_dotted_order = ( + generate_dotted_order(trace_info.message_id, trace_info.start_time) if trace_info.message_id else None + ) + workflow_dotted_order = generate_dotted_order( + trace_info.workflow_run_id, + trace_info.workflow_data.created_at, + message_dotted_order, + ) + metadata = trace_info.metadata + metadata["workflow_app_log_id"] = trace_info.workflow_app_log_id + + if trace_info.message_id: + message_run = LangSmithRunModel( + id=trace_info.message_id, + name=TraceTaskName.MESSAGE_TRACE.value, + inputs=dict(trace_info.workflow_run_inputs), + outputs=dict(trace_info.workflow_run_outputs), + run_type=LangSmithRunType.chain, + start_time=trace_info.start_time, + end_time=trace_info.end_time, + extra={ + "metadata": metadata, + }, + tags=["message", "workflow"], + error=trace_info.error, + trace_id=trace_id, + dotted_order=message_dotted_order, + file_list=[], + serialized=None, + parent_run_id=None, + events=[], + session_id=None, + session_name=None, + reference_example_id=None, + input_attachments={}, + output_attachments={}, + ) + self.add_run(message_run) + + langsmith_run = LangSmithRunModel( + file_list=trace_info.file_list, + total_tokens=trace_info.total_tokens, + id=trace_info.workflow_run_id, + name=TraceTaskName.WORKFLOW_TRACE.value, + inputs=dict(trace_info.workflow_run_inputs), + run_type=LangSmithRunType.tool, + start_time=trace_info.workflow_data.created_at, + end_time=trace_info.workflow_data.finished_at, + outputs=dict(trace_info.workflow_run_outputs), + extra={ + "metadata": metadata, + }, + error=trace_info.error, + tags=["workflow"], + parent_run_id=trace_info.message_id or None, + trace_id=trace_id, + dotted_order=workflow_dotted_order, + serialized=None, + events=[], + session_id=None, + session_name=None, + reference_example_id=None, + input_attachments={}, + output_attachments={}, + ) + + self.add_run(langsmith_run) + + # through workflow_run_id get all_nodes_execution + workflow_nodes_execution_id_records = ( + db.session.query(WorkflowNodeExecution.id) + .filter(WorkflowNodeExecution.workflow_run_id == trace_info.workflow_run_id) + .all() + ) + + for node_execution_id_record in workflow_nodes_execution_id_records: + node_execution = ( + db.session.query( + WorkflowNodeExecution.id, + WorkflowNodeExecution.tenant_id, + WorkflowNodeExecution.app_id, + WorkflowNodeExecution.title, + WorkflowNodeExecution.node_type, + WorkflowNodeExecution.status, + WorkflowNodeExecution.inputs, + WorkflowNodeExecution.outputs, + WorkflowNodeExecution.created_at, + WorkflowNodeExecution.elapsed_time, + WorkflowNodeExecution.process_data, + WorkflowNodeExecution.execution_metadata, + ) + .filter(WorkflowNodeExecution.id == node_execution_id_record.id) + .first() + ) + + if not node_execution: + continue + + node_execution_id = node_execution.id + tenant_id = node_execution.tenant_id + app_id = node_execution.app_id + node_name = node_execution.title + node_type = node_execution.node_type + status = node_execution.status + if node_type == "llm": + inputs = ( + json.loads(node_execution.process_data).get("prompts", {}) if node_execution.process_data else {} + ) + else: + inputs = json.loads(node_execution.inputs) if node_execution.inputs else {} + outputs = json.loads(node_execution.outputs) if node_execution.outputs else {} + created_at = node_execution.created_at or datetime.now() + elapsed_time = node_execution.elapsed_time + finished_at = created_at + timedelta(seconds=elapsed_time) + + execution_metadata = ( + json.loads(node_execution.execution_metadata) if node_execution.execution_metadata else {} + ) + node_total_tokens = execution_metadata.get("total_tokens", 0) + metadata = execution_metadata.copy() + metadata.update( + { + "workflow_run_id": trace_info.workflow_run_id, + "node_execution_id": node_execution_id, + "tenant_id": tenant_id, + "app_id": app_id, + "app_name": node_name, + "node_type": node_type, + "status": status, + } + ) + + process_data = json.loads(node_execution.process_data) if node_execution.process_data else {} + if process_data and process_data.get("model_mode") == "chat": + run_type = LangSmithRunType.llm + metadata.update( + { + "ls_provider": process_data.get("model_provider", ""), + "ls_model_name": process_data.get("model_name", ""), + } + ) + elif node_type == "knowledge-retrieval": + run_type = LangSmithRunType.retriever + else: + run_type = LangSmithRunType.tool + + node_dotted_order = generate_dotted_order(node_execution_id, created_at, workflow_dotted_order) + langsmith_run = LangSmithRunModel( + total_tokens=node_total_tokens, + name=node_type, + inputs=inputs, + run_type=run_type, + start_time=created_at, + end_time=finished_at, + outputs=outputs, + file_list=trace_info.file_list, + extra={ + "metadata": metadata, + }, + parent_run_id=trace_info.workflow_run_id, + tags=["node_execution"], + id=node_execution_id, + trace_id=trace_id, + dotted_order=node_dotted_order, + error="", + serialized=None, + events=[], + session_id=None, + session_name=None, + reference_example_id=None, + input_attachments={}, + output_attachments={}, + ) + + self.add_run(langsmith_run) + + def message_trace(self, trace_info: MessageTraceInfo): + # get message file data + file_list = cast(list[str], trace_info.file_list) or [] + message_file_data: Optional[MessageFile] = trace_info.message_file_data + file_url = f"{self.file_base_url}/{message_file_data.url}" if message_file_data else "" + file_list.append(file_url) + metadata = trace_info.metadata + message_data = trace_info.message_data + if message_data is None: + return + message_id = message_data.id + + user_id = message_data.from_account_id + metadata["user_id"] = user_id + + if message_data.from_end_user_id: + end_user_data: Optional[EndUser] = ( + db.session.query(EndUser).filter(EndUser.id == message_data.from_end_user_id).first() + ) + if end_user_data is not None: + end_user_id = end_user_data.session_id + metadata["end_user_id"] = end_user_id + + message_run = LangSmithRunModel( + input_tokens=trace_info.message_tokens, + output_tokens=trace_info.answer_tokens, + total_tokens=trace_info.total_tokens, + id=message_id, + name=TraceTaskName.MESSAGE_TRACE.value, + inputs=trace_info.inputs, + run_type=LangSmithRunType.chain, + start_time=trace_info.start_time, + end_time=trace_info.end_time, + outputs=message_data.answer, + extra={"metadata": metadata}, + tags=["message", str(trace_info.conversation_mode)], + error=trace_info.error, + file_list=file_list, + serialized=None, + events=[], + session_id=None, + session_name=None, + reference_example_id=None, + input_attachments={}, + output_attachments={}, + trace_id=None, + dotted_order=None, + parent_run_id=None, + ) + self.add_run(message_run) + + # create llm run parented to message run + llm_run = LangSmithRunModel( + input_tokens=trace_info.message_tokens, + output_tokens=trace_info.answer_tokens, + total_tokens=trace_info.total_tokens, + name="llm", + inputs=trace_info.inputs, + run_type=LangSmithRunType.llm, + start_time=trace_info.start_time, + end_time=trace_info.end_time, + outputs=message_data.answer, + extra={"metadata": metadata}, + parent_run_id=message_id, + tags=["llm", str(trace_info.conversation_mode)], + error=trace_info.error, + file_list=file_list, + serialized=None, + events=[], + session_id=None, + session_name=None, + reference_example_id=None, + input_attachments={}, + output_attachments={}, + trace_id=None, + dotted_order=None, + id=str(uuid.uuid4()), + ) + self.add_run(llm_run) + + def moderation_trace(self, trace_info: ModerationTraceInfo): + if trace_info.message_data is None: + return + langsmith_run = LangSmithRunModel( + name=TraceTaskName.MODERATION_TRACE.value, + inputs=trace_info.inputs, + outputs={ + "action": trace_info.action, + "flagged": trace_info.flagged, + "preset_response": trace_info.preset_response, + "inputs": trace_info.inputs, + }, + run_type=LangSmithRunType.tool, + extra={"metadata": trace_info.metadata}, + tags=["moderation"], + parent_run_id=trace_info.message_id, + start_time=trace_info.start_time or trace_info.message_data.created_at, + end_time=trace_info.end_time or trace_info.message_data.updated_at, + id=str(uuid.uuid4()), + serialized=None, + events=[], + session_id=None, + session_name=None, + reference_example_id=None, + input_attachments={}, + output_attachments={}, + trace_id=None, + dotted_order=None, + error="", + file_list=[], + ) + + self.add_run(langsmith_run) + + def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo): + message_data = trace_info.message_data + if message_data is None: + return + suggested_question_run = LangSmithRunModel( + name=TraceTaskName.SUGGESTED_QUESTION_TRACE.value, + inputs=trace_info.inputs, + outputs=trace_info.suggested_question, + run_type=LangSmithRunType.tool, + extra={"metadata": trace_info.metadata}, + tags=["suggested_question"], + parent_run_id=trace_info.message_id, + start_time=trace_info.start_time or message_data.created_at, + end_time=trace_info.end_time or message_data.updated_at, + id=str(uuid.uuid4()), + serialized=None, + events=[], + session_id=None, + session_name=None, + reference_example_id=None, + input_attachments={}, + output_attachments={}, + trace_id=None, + dotted_order=None, + error="", + file_list=[], + ) + + self.add_run(suggested_question_run) + + def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo): + if trace_info.message_data is None: + return + dataset_retrieval_run = LangSmithRunModel( + name=TraceTaskName.DATASET_RETRIEVAL_TRACE.value, + inputs=trace_info.inputs, + outputs={"documents": trace_info.documents}, + run_type=LangSmithRunType.retriever, + extra={"metadata": trace_info.metadata}, + tags=["dataset_retrieval"], + parent_run_id=trace_info.message_id, + start_time=trace_info.start_time or trace_info.message_data.created_at, + end_time=trace_info.end_time or trace_info.message_data.updated_at, + id=str(uuid.uuid4()), + serialized=None, + events=[], + session_id=None, + session_name=None, + reference_example_id=None, + input_attachments={}, + output_attachments={}, + trace_id=None, + dotted_order=None, + error="", + file_list=[], + ) + + self.add_run(dataset_retrieval_run) + + def tool_trace(self, trace_info: ToolTraceInfo): + tool_run = LangSmithRunModel( + name=trace_info.tool_name, + inputs=trace_info.tool_inputs, + outputs=trace_info.tool_outputs, + run_type=LangSmithRunType.tool, + extra={ + "metadata": trace_info.metadata, + }, + tags=["tool", trace_info.tool_name], + parent_run_id=trace_info.message_id, + start_time=trace_info.start_time, + end_time=trace_info.end_time, + file_list=[cast(str, trace_info.file_url)], + id=str(uuid.uuid4()), + serialized=None, + events=[], + session_id=None, + session_name=None, + reference_example_id=None, + input_attachments={}, + output_attachments={}, + trace_id=None, + dotted_order=None, + error=trace_info.error or "", + ) + + self.add_run(tool_run) + + def generate_name_trace(self, trace_info: GenerateNameTraceInfo): + name_run = LangSmithRunModel( + name=TraceTaskName.GENERATE_NAME_TRACE.value, + inputs=trace_info.inputs, + outputs=trace_info.outputs, + run_type=LangSmithRunType.tool, + extra={"metadata": trace_info.metadata}, + tags=["generate_name"], + start_time=trace_info.start_time or datetime.now(), + end_time=trace_info.end_time or datetime.now(), + id=str(uuid.uuid4()), + serialized=None, + events=[], + session_id=None, + session_name=None, + reference_example_id=None, + input_attachments={}, + output_attachments={}, + trace_id=None, + dotted_order=None, + error="", + file_list=[], + parent_run_id=None, + ) + + self.add_run(name_run) + + def add_run(self, run_data: LangSmithRunModel): + data = run_data.model_dump() + if self.project_id: + data["session_id"] = self.project_id + elif self.project_name: + data["session_name"] = self.project_name + + data = filter_none_values(data) + try: + self.langsmith_client.create_run(**data) + logger.debug("LangSmith Run created successfully.") + except Exception as e: + raise ValueError(f"LangSmith Failed to create run: {str(e)}") + + def update_run(self, update_run_data: LangSmithRunUpdateModel): + data = update_run_data.model_dump() + data = filter_none_values(data) + try: + self.langsmith_client.update_run(**data) + logger.debug("LangSmith Run updated successfully.") + except Exception as e: + raise ValueError(f"LangSmith Failed to update run: {str(e)}") + + def api_check(self): + try: + random_project_name = f"test_project_{datetime.now().strftime('%Y%m%d%H%M%S')}" + self.langsmith_client.create_project(project_name=random_project_name) + self.langsmith_client.delete_project(project_name=random_project_name) + return True + except Exception as e: + logger.debug(f"LangSmith API check failed: {str(e)}") + raise ValueError(f"LangSmith API check failed: {str(e)}") + + def get_project_url(self): + try: + run_data = RunBase( + id=uuid.uuid4(), + name="tool", + inputs={"input": "test"}, + outputs={"output": "test"}, + run_type=LangSmithRunType.tool, + start_time=datetime.now(), + ) + + project_url = self.langsmith_client.get_run_url( + run=run_data, project_id=self.project_id, project_name=self.project_name + ) + return project_url.split("/r/")[0] + except Exception as e: + logger.debug(f"LangSmith get run url failed: {str(e)}") + raise ValueError(f"LangSmith get run url failed: {str(e)}") diff --git a/api/core/ops/opik_trace/__init__.py b/api/core/ops/opik_trace/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/ops/opik_trace/opik_trace.py b/api/core/ops/opik_trace/opik_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..fabf38fbd641786140ea38677c801a7caf0b414a --- /dev/null +++ b/api/core/ops/opik_trace/opik_trace.py @@ -0,0 +1,469 @@ +import json +import logging +import os +import uuid +from datetime import datetime, timedelta +from typing import Optional, cast + +from opik import Opik, Trace +from opik.id_helpers import uuid4_to_uuid7 + +from core.ops.base_trace_instance import BaseTraceInstance +from core.ops.entities.config_entity import OpikConfig +from core.ops.entities.trace_entity import ( + BaseTraceInfo, + DatasetRetrievalTraceInfo, + GenerateNameTraceInfo, + MessageTraceInfo, + ModerationTraceInfo, + SuggestedQuestionTraceInfo, + ToolTraceInfo, + TraceTaskName, + WorkflowTraceInfo, +) +from extensions.ext_database import db +from models.model import EndUser, MessageFile +from models.workflow import WorkflowNodeExecution + +logger = logging.getLogger(__name__) + + +def wrap_dict(key_name, data): + """Make sure that the input data is a dict""" + if not isinstance(data, dict): + return {key_name: data} + + return data + + +def wrap_metadata(metadata, **kwargs): + """Add common metatada to all Traces and Spans""" + metadata["created_from"] = "dify" + + metadata.update(kwargs) + + return metadata + + +def prepare_opik_uuid(user_datetime: Optional[datetime], user_uuid: Optional[str]): + """Opik needs UUIDv7 while Dify uses UUIDv4 for identifier of most + messages and objects. The type-hints of BaseTraceInfo indicates that + objects start_time and message_id could be null which means we cannot map + it to a UUIDv7. Given that we have no way to identify that object + uniquely, generate a new random one UUIDv7 in that case. + """ + + if user_datetime is None: + user_datetime = datetime.now() + + if user_uuid is None: + user_uuid = str(uuid.uuid4()) + + return uuid4_to_uuid7(user_datetime, user_uuid) + + +class OpikDataTrace(BaseTraceInstance): + def __init__( + self, + opik_config: OpikConfig, + ): + super().__init__(opik_config) + self.opik_client = Opik( + project_name=opik_config.project, + workspace=opik_config.workspace, + host=opik_config.url, + api_key=opik_config.api_key, + ) + self.project = opik_config.project + self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001") + + def trace(self, trace_info: BaseTraceInfo): + if isinstance(trace_info, WorkflowTraceInfo): + self.workflow_trace(trace_info) + if isinstance(trace_info, MessageTraceInfo): + self.message_trace(trace_info) + if isinstance(trace_info, ModerationTraceInfo): + self.moderation_trace(trace_info) + if isinstance(trace_info, SuggestedQuestionTraceInfo): + self.suggested_question_trace(trace_info) + if isinstance(trace_info, DatasetRetrievalTraceInfo): + self.dataset_retrieval_trace(trace_info) + if isinstance(trace_info, ToolTraceInfo): + self.tool_trace(trace_info) + if isinstance(trace_info, GenerateNameTraceInfo): + self.generate_name_trace(trace_info) + + def workflow_trace(self, trace_info: WorkflowTraceInfo): + dify_trace_id = trace_info.workflow_run_id + opik_trace_id = prepare_opik_uuid(trace_info.start_time, dify_trace_id) + workflow_metadata = wrap_metadata( + trace_info.metadata, message_id=trace_info.message_id, workflow_app_log_id=trace_info.workflow_app_log_id + ) + root_span_id = None + + if trace_info.message_id: + dify_trace_id = trace_info.message_id + opik_trace_id = prepare_opik_uuid(trace_info.start_time, dify_trace_id) + + trace_data = { + "id": opik_trace_id, + "name": TraceTaskName.MESSAGE_TRACE.value, + "start_time": trace_info.start_time, + "end_time": trace_info.end_time, + "metadata": workflow_metadata, + "input": wrap_dict("input", trace_info.workflow_run_inputs), + "output": wrap_dict("output", trace_info.workflow_run_outputs), + "tags": ["message", "workflow"], + "project_name": self.project, + } + self.add_trace(trace_data) + + root_span_id = prepare_opik_uuid(trace_info.start_time, trace_info.workflow_run_id) + span_data = { + "id": root_span_id, + "parent_span_id": None, + "trace_id": opik_trace_id, + "name": TraceTaskName.WORKFLOW_TRACE.value, + "input": wrap_dict("input", trace_info.workflow_run_inputs), + "output": wrap_dict("output", trace_info.workflow_run_outputs), + "start_time": trace_info.start_time, + "end_time": trace_info.end_time, + "metadata": workflow_metadata, + "tags": ["workflow"], + "project_name": self.project, + } + self.add_span(span_data) + else: + trace_data = { + "id": opik_trace_id, + "name": TraceTaskName.MESSAGE_TRACE.value, + "start_time": trace_info.start_time, + "end_time": trace_info.end_time, + "metadata": workflow_metadata, + "input": wrap_dict("input", trace_info.workflow_run_inputs), + "output": wrap_dict("output", trace_info.workflow_run_outputs), + "tags": ["workflow"], + "project_name": self.project, + } + self.add_trace(trace_data) + + # through workflow_run_id get all_nodes_execution + workflow_nodes_execution_id_records = ( + db.session.query(WorkflowNodeExecution.id) + .filter(WorkflowNodeExecution.workflow_run_id == trace_info.workflow_run_id) + .all() + ) + + for node_execution_id_record in workflow_nodes_execution_id_records: + node_execution = ( + db.session.query( + WorkflowNodeExecution.id, + WorkflowNodeExecution.tenant_id, + WorkflowNodeExecution.app_id, + WorkflowNodeExecution.title, + WorkflowNodeExecution.node_type, + WorkflowNodeExecution.status, + WorkflowNodeExecution.inputs, + WorkflowNodeExecution.outputs, + WorkflowNodeExecution.created_at, + WorkflowNodeExecution.elapsed_time, + WorkflowNodeExecution.process_data, + WorkflowNodeExecution.execution_metadata, + ) + .filter(WorkflowNodeExecution.id == node_execution_id_record.id) + .first() + ) + + if not node_execution: + continue + + node_execution_id = node_execution.id + tenant_id = node_execution.tenant_id + app_id = node_execution.app_id + node_name = node_execution.title + node_type = node_execution.node_type + status = node_execution.status + if node_type == "llm": + inputs = ( + json.loads(node_execution.process_data).get("prompts", {}) if node_execution.process_data else {} + ) + else: + inputs = json.loads(node_execution.inputs) if node_execution.inputs else {} + outputs = json.loads(node_execution.outputs) if node_execution.outputs else {} + created_at = node_execution.created_at or datetime.now() + elapsed_time = node_execution.elapsed_time + finished_at = created_at + timedelta(seconds=elapsed_time) + + execution_metadata = ( + json.loads(node_execution.execution_metadata) if node_execution.execution_metadata else {} + ) + metadata = execution_metadata.copy() + metadata.update( + { + "workflow_run_id": trace_info.workflow_run_id, + "node_execution_id": node_execution_id, + "tenant_id": tenant_id, + "app_id": app_id, + "app_name": node_name, + "node_type": node_type, + "status": status, + } + ) + + process_data = json.loads(node_execution.process_data) if node_execution.process_data else {} + + provider = None + model = None + total_tokens = 0 + completion_tokens = 0 + prompt_tokens = 0 + + if process_data and process_data.get("model_mode") == "chat": + run_type = "llm" + provider = process_data.get("model_provider", None) + model = process_data.get("model_name", "") + metadata.update( + { + "ls_provider": provider, + "ls_model_name": model, + } + ) + + try: + if outputs.get("usage"): + total_tokens = outputs["usage"].get("total_tokens", 0) + prompt_tokens = outputs["usage"].get("prompt_tokens", 0) + completion_tokens = outputs["usage"].get("completion_tokens", 0) + except Exception: + logger.error("Failed to extract usage", exc_info=True) + + else: + run_type = "tool" + + parent_span_id = trace_info.workflow_app_log_id or trace_info.workflow_run_id + + if not total_tokens: + total_tokens = execution_metadata.get("total_tokens", 0) + + span_data = { + "trace_id": opik_trace_id, + "id": prepare_opik_uuid(created_at, node_execution_id), + "parent_span_id": prepare_opik_uuid(trace_info.start_time, parent_span_id), + "name": node_type, + "type": run_type, + "start_time": created_at, + "end_time": finished_at, + "metadata": wrap_metadata(metadata), + "input": wrap_dict("input", inputs), + "output": wrap_dict("output", outputs), + "tags": ["node_execution"], + "project_name": self.project, + "usage": { + "total_tokens": total_tokens, + "completion_tokens": completion_tokens, + "prompt_tokens": prompt_tokens, + }, + "model": model, + "provider": provider, + } + + self.add_span(span_data) + + def message_trace(self, trace_info: MessageTraceInfo): + # get message file data + file_list = cast(list[str], trace_info.file_list) or [] + message_file_data: Optional[MessageFile] = trace_info.message_file_data + + if message_file_data is not None: + file_url = f"{self.file_base_url}/{message_file_data.url}" if message_file_data else "" + file_list.append(file_url) + + message_data = trace_info.message_data + if message_data is None: + return + + metadata = trace_info.metadata + message_id = trace_info.message_id + + user_id = message_data.from_account_id + metadata["user_id"] = user_id + metadata["file_list"] = file_list + + if message_data.from_end_user_id: + end_user_data: Optional[EndUser] = ( + db.session.query(EndUser).filter(EndUser.id == message_data.from_end_user_id).first() + ) + if end_user_data is not None: + end_user_id = end_user_data.session_id + metadata["end_user_id"] = end_user_id + + trace_data = { + "id": prepare_opik_uuid(trace_info.start_time, message_id), + "name": TraceTaskName.MESSAGE_TRACE.value, + "start_time": trace_info.start_time, + "end_time": trace_info.end_time, + "metadata": wrap_metadata(metadata), + "input": trace_info.inputs, + "output": message_data.answer, + "tags": ["message", str(trace_info.conversation_mode)], + "project_name": self.project, + } + trace = self.add_trace(trace_data) + + span_data = { + "trace_id": trace.id, + "name": "llm", + "type": "llm", + "start_time": trace_info.start_time, + "end_time": trace_info.end_time, + "metadata": wrap_metadata(metadata), + "input": {"input": trace_info.inputs}, + "output": {"output": message_data.answer}, + "tags": ["llm", str(trace_info.conversation_mode)], + "usage": { + "completion_tokens": trace_info.answer_tokens, + "prompt_tokens": trace_info.message_tokens, + "total_tokens": trace_info.total_tokens, + }, + "project_name": self.project, + } + self.add_span(span_data) + + def moderation_trace(self, trace_info: ModerationTraceInfo): + if trace_info.message_data is None: + return + + start_time = trace_info.start_time or trace_info.message_data.created_at + + span_data = { + "trace_id": prepare_opik_uuid(start_time, trace_info.message_id), + "name": TraceTaskName.MODERATION_TRACE.value, + "type": "tool", + "start_time": start_time, + "end_time": trace_info.end_time or trace_info.message_data.updated_at, + "metadata": wrap_metadata(trace_info.metadata), + "input": wrap_dict("input", trace_info.inputs), + "output": { + "action": trace_info.action, + "flagged": trace_info.flagged, + "preset_response": trace_info.preset_response, + "inputs": trace_info.inputs, + }, + "tags": ["moderation"], + } + + self.add_span(span_data) + + def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo): + message_data = trace_info.message_data + if message_data is None: + return + + start_time = trace_info.start_time or message_data.created_at + + span_data = { + "trace_id": prepare_opik_uuid(start_time, trace_info.message_id), + "name": TraceTaskName.SUGGESTED_QUESTION_TRACE.value, + "type": "tool", + "start_time": start_time, + "end_time": trace_info.end_time or message_data.updated_at, + "metadata": wrap_metadata(trace_info.metadata), + "input": wrap_dict("input", trace_info.inputs), + "output": wrap_dict("output", trace_info.suggested_question), + "tags": ["suggested_question"], + } + + self.add_span(span_data) + + def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo): + if trace_info.message_data is None: + return + + start_time = trace_info.start_time or trace_info.message_data.created_at + + span_data = { + "trace_id": prepare_opik_uuid(start_time, trace_info.message_id), + "name": TraceTaskName.DATASET_RETRIEVAL_TRACE.value, + "type": "tool", + "start_time": start_time, + "end_time": trace_info.end_time or trace_info.message_data.updated_at, + "metadata": wrap_metadata(trace_info.metadata), + "input": wrap_dict("input", trace_info.inputs), + "output": {"documents": trace_info.documents}, + "tags": ["dataset_retrieval"], + } + + self.add_span(span_data) + + def tool_trace(self, trace_info: ToolTraceInfo): + span_data = { + "trace_id": prepare_opik_uuid(trace_info.start_time, trace_info.message_id), + "name": trace_info.tool_name, + "type": "tool", + "start_time": trace_info.start_time, + "end_time": trace_info.end_time, + "metadata": wrap_metadata(trace_info.metadata), + "input": wrap_dict("input", trace_info.tool_inputs), + "output": wrap_dict("output", trace_info.tool_outputs), + "tags": ["tool", trace_info.tool_name], + } + + self.add_span(span_data) + + def generate_name_trace(self, trace_info: GenerateNameTraceInfo): + trace_data = { + "id": prepare_opik_uuid(trace_info.start_time, trace_info.message_id), + "name": TraceTaskName.GENERATE_NAME_TRACE.value, + "start_time": trace_info.start_time, + "end_time": trace_info.end_time, + "metadata": wrap_metadata(trace_info.metadata), + "input": trace_info.inputs, + "output": trace_info.outputs, + "tags": ["generate_name"], + "project_name": self.project, + } + + trace = self.add_trace(trace_data) + + span_data = { + "trace_id": trace.id, + "name": TraceTaskName.GENERATE_NAME_TRACE.value, + "start_time": trace_info.start_time, + "end_time": trace_info.end_time, + "metadata": wrap_metadata(trace_info.metadata), + "input": wrap_dict("input", trace_info.inputs), + "output": wrap_dict("output", trace_info.outputs), + "tags": ["generate_name"], + } + + self.add_span(span_data) + + def add_trace(self, opik_trace_data: dict) -> Trace: + try: + trace = self.opik_client.trace(**opik_trace_data) + logger.debug("Opik Trace created successfully") + return trace + except Exception as e: + raise ValueError(f"Opik Failed to create trace: {str(e)}") + + def add_span(self, opik_span_data: dict): + try: + self.opik_client.span(**opik_span_data) + logger.debug("Opik Span created successfully") + except Exception as e: + raise ValueError(f"Opik Failed to create span: {str(e)}") + + def api_check(self): + try: + self.opik_client.auth_check() + return True + except Exception as e: + logger.info(f"Opik API check failed: {str(e)}", exc_info=True) + raise ValueError(f"Opik API check failed: {str(e)}") + + def get_project_url(self): + try: + return self.opik_client.get_project_url(project_name=self.project) + except Exception as e: + logger.info(f"Opik get run url failed: {str(e)}", exc_info=True) + raise ValueError(f"Opik get run url failed: {str(e)}") diff --git a/api/core/ops/ops_trace_manager.py b/api/core/ops/ops_trace_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..c153e3f9dd6d47e06c45f0df94b3cdfce4e19270 --- /dev/null +++ b/api/core/ops/ops_trace_manager.py @@ -0,0 +1,811 @@ +import json +import logging +import os +import queue +import threading +import time +from datetime import timedelta +from typing import Any, Optional, Union +from uuid import UUID, uuid4 + +from flask import current_app +from sqlalchemy import select +from sqlalchemy.orm import Session + +from core.helper.encrypter import decrypt_token, encrypt_token, obfuscated_token +from core.ops.entities.config_entity import ( + OPS_FILE_PATH, + LangfuseConfig, + LangSmithConfig, + OpikConfig, + TracingProviderEnum, +) +from core.ops.entities.trace_entity import ( + DatasetRetrievalTraceInfo, + GenerateNameTraceInfo, + MessageTraceInfo, + ModerationTraceInfo, + SuggestedQuestionTraceInfo, + TaskData, + ToolTraceInfo, + TraceTaskName, + WorkflowTraceInfo, +) +from core.ops.langfuse_trace.langfuse_trace import LangFuseDataTrace +from core.ops.langsmith_trace.langsmith_trace import LangSmithDataTrace +from core.ops.opik_trace.opik_trace import OpikDataTrace +from core.ops.utils import get_message_data +from extensions.ext_database import db +from extensions.ext_storage import storage +from models.model import App, AppModelConfig, Conversation, Message, MessageFile, TraceAppConfig +from models.workflow import WorkflowAppLog, WorkflowRun +from tasks.ops_trace_task import process_trace_tasks + +provider_config_map: dict[str, dict[str, Any]] = { + TracingProviderEnum.LANGFUSE.value: { + "config_class": LangfuseConfig, + "secret_keys": ["public_key", "secret_key"], + "other_keys": ["host", "project_key"], + "trace_instance": LangFuseDataTrace, + }, + TracingProviderEnum.LANGSMITH.value: { + "config_class": LangSmithConfig, + "secret_keys": ["api_key"], + "other_keys": ["project", "endpoint"], + "trace_instance": LangSmithDataTrace, + }, + TracingProviderEnum.OPIK.value: { + "config_class": OpikConfig, + "secret_keys": ["api_key"], + "other_keys": ["project", "url", "workspace"], + "trace_instance": OpikDataTrace, + }, +} + + +class OpsTraceManager: + @classmethod + def encrypt_tracing_config( + cls, tenant_id: str, tracing_provider: str, tracing_config: dict, current_trace_config=None + ): + """ + Encrypt tracing config. + :param tenant_id: tenant id + :param tracing_provider: tracing provider + :param tracing_config: tracing config dictionary to be encrypted + :param current_trace_config: current tracing configuration for keeping existing values + :return: encrypted tracing configuration + """ + # Get the configuration class and the keys that require encryption + config_class, secret_keys, other_keys = ( + provider_config_map[tracing_provider]["config_class"], + provider_config_map[tracing_provider]["secret_keys"], + provider_config_map[tracing_provider]["other_keys"], + ) + + new_config = {} + # Encrypt necessary keys + for key in secret_keys: + if key in tracing_config: + if "*" in tracing_config[key]: + # If the key contains '*', retain the original value from the current config + new_config[key] = current_trace_config.get(key, tracing_config[key]) + else: + # Otherwise, encrypt the key + new_config[key] = encrypt_token(tenant_id, tracing_config[key]) + + for key in other_keys: + new_config[key] = tracing_config.get(key, "") + + # Create a new instance of the config class with the new configuration + encrypted_config = config_class(**new_config) + return encrypted_config.model_dump() + + @classmethod + def decrypt_tracing_config(cls, tenant_id: str, tracing_provider: str, tracing_config: dict): + """ + Decrypt tracing config + :param tenant_id: tenant id + :param tracing_provider: tracing provider + :param tracing_config: tracing config + :return: + """ + config_class, secret_keys, other_keys = ( + provider_config_map[tracing_provider]["config_class"], + provider_config_map[tracing_provider]["secret_keys"], + provider_config_map[tracing_provider]["other_keys"], + ) + new_config = {} + for key in secret_keys: + if key in tracing_config: + new_config[key] = decrypt_token(tenant_id, tracing_config[key]) + + for key in other_keys: + new_config[key] = tracing_config.get(key, "") + + return config_class(**new_config).model_dump() + + @classmethod + def obfuscated_decrypt_token(cls, tracing_provider: str, decrypt_tracing_config: dict): + """ + Decrypt tracing config + :param tracing_provider: tracing provider + :param decrypt_tracing_config: tracing config + :return: + """ + config_class, secret_keys, other_keys = ( + provider_config_map[tracing_provider]["config_class"], + provider_config_map[tracing_provider]["secret_keys"], + provider_config_map[tracing_provider]["other_keys"], + ) + new_config = {} + for key in secret_keys: + if key in decrypt_tracing_config: + new_config[key] = obfuscated_token(decrypt_tracing_config[key]) + + for key in other_keys: + new_config[key] = decrypt_tracing_config.get(key, "") + return config_class(**new_config).model_dump() + + @classmethod + def get_decrypted_tracing_config(cls, app_id: str, tracing_provider: str): + """ + Get decrypted tracing config + :param app_id: app id + :param tracing_provider: tracing provider + :return: + """ + trace_config_data: Optional[TraceAppConfig] = ( + db.session.query(TraceAppConfig) + .filter(TraceAppConfig.app_id == app_id, TraceAppConfig.tracing_provider == tracing_provider) + .first() + ) + + if not trace_config_data: + return None + + # decrypt_token + app = db.session.query(App).filter(App.id == app_id).first() + if not app: + raise ValueError("App not found") + + tenant_id = app.tenant_id + decrypt_tracing_config = cls.decrypt_tracing_config( + tenant_id, tracing_provider, trace_config_data.tracing_config + ) + + return decrypt_tracing_config + + @classmethod + def get_ops_trace_instance( + cls, + app_id: Optional[Union[UUID, str]] = None, + ): + """ + Get ops trace through model config + :param app_id: app_id + :return: + """ + if isinstance(app_id, UUID): + app_id = str(app_id) + + if app_id is None: + return None + + app: Optional[App] = db.session.query(App).filter(App.id == app_id).first() + + if app is None: + return None + + app_ops_trace_config = json.loads(app.tracing) if app.tracing else None + + if app_ops_trace_config is None: + return None + + tracing_provider = app_ops_trace_config.get("tracing_provider") + + if tracing_provider is None or tracing_provider not in provider_config_map: + return None + + # decrypt_token + decrypt_trace_config = cls.get_decrypted_tracing_config(app_id, tracing_provider) + if app_ops_trace_config.get("enabled"): + trace_instance, config_class = ( + provider_config_map[tracing_provider]["trace_instance"], + provider_config_map[tracing_provider]["config_class"], + ) + tracing_instance = trace_instance(config_class(**decrypt_trace_config)) + return tracing_instance + + return None + + @classmethod + def get_app_config_through_message_id(cls, message_id: str): + app_model_config = None + message_data = db.session.query(Message).filter(Message.id == message_id).first() + if not message_data: + return None + conversation_id = message_data.conversation_id + conversation_data = db.session.query(Conversation).filter(Conversation.id == conversation_id).first() + if not conversation_data: + return None + + if conversation_data.app_model_config_id: + app_model_config = ( + db.session.query(AppModelConfig) + .filter(AppModelConfig.id == conversation_data.app_model_config_id) + .first() + ) + elif conversation_data.app_model_config_id is None and conversation_data.override_model_configs: + app_model_config = conversation_data.override_model_configs + + return app_model_config + + @classmethod + def update_app_tracing_config(cls, app_id: str, enabled: bool, tracing_provider: str): + """ + Update app tracing config + :param app_id: app id + :param enabled: enabled + :param tracing_provider: tracing provider + :return: + """ + # auth check + if tracing_provider not in provider_config_map and tracing_provider is not None: + raise ValueError(f"Invalid tracing provider: {tracing_provider}") + + app_config: Optional[App] = db.session.query(App).filter(App.id == app_id).first() + if not app_config: + raise ValueError("App not found") + app_config.tracing = json.dumps( + { + "enabled": enabled, + "tracing_provider": tracing_provider, + } + ) + db.session.commit() + + @classmethod + def get_app_tracing_config(cls, app_id: str): + """ + Get app tracing config + :param app_id: app id + :return: + """ + app: Optional[App] = db.session.query(App).filter(App.id == app_id).first() + if not app: + raise ValueError("App not found") + if not app.tracing: + return {"enabled": False, "tracing_provider": None} + app_trace_config = json.loads(app.tracing) + return app_trace_config + + @staticmethod + def check_trace_config_is_effective(tracing_config: dict, tracing_provider: str): + """ + Check trace config is effective + :param tracing_config: tracing config + :param tracing_provider: tracing provider + :return: + """ + config_type, trace_instance = ( + provider_config_map[tracing_provider]["config_class"], + provider_config_map[tracing_provider]["trace_instance"], + ) + tracing_config = config_type(**tracing_config) + return trace_instance(tracing_config).api_check() + + @staticmethod + def get_trace_config_project_key(tracing_config: dict, tracing_provider: str): + """ + get trace config is project key + :param tracing_config: tracing config + :param tracing_provider: tracing provider + :return: + """ + config_type, trace_instance = ( + provider_config_map[tracing_provider]["config_class"], + provider_config_map[tracing_provider]["trace_instance"], + ) + tracing_config = config_type(**tracing_config) + return trace_instance(tracing_config).get_project_key() + + @staticmethod + def get_trace_config_project_url(tracing_config: dict, tracing_provider: str): + """ + get trace config is project key + :param tracing_config: tracing config + :param tracing_provider: tracing provider + :return: + """ + config_type, trace_instance = ( + provider_config_map[tracing_provider]["config_class"], + provider_config_map[tracing_provider]["trace_instance"], + ) + tracing_config = config_type(**tracing_config) + return trace_instance(tracing_config).get_project_url() + + +class TraceTask: + def __init__( + self, + trace_type: Any, + message_id: Optional[str] = None, + workflow_run: Optional[WorkflowRun] = None, + conversation_id: Optional[str] = None, + user_id: Optional[str] = None, + timer: Optional[Any] = None, + **kwargs, + ): + self.trace_type = trace_type + self.message_id = message_id + self.workflow_run_id = workflow_run.id if workflow_run else None + self.conversation_id = conversation_id + self.user_id = user_id + self.timer = timer + self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001") + self.app_id = None + + self.kwargs = kwargs + + def execute(self): + return self.preprocess() + + def preprocess(self): + preprocess_map = { + TraceTaskName.CONVERSATION_TRACE: lambda: self.conversation_trace(**self.kwargs), + TraceTaskName.WORKFLOW_TRACE: lambda: self.workflow_trace( + workflow_run_id=self.workflow_run_id, conversation_id=self.conversation_id, user_id=self.user_id + ), + TraceTaskName.MESSAGE_TRACE: lambda: self.message_trace(message_id=self.message_id), + TraceTaskName.MODERATION_TRACE: lambda: self.moderation_trace( + message_id=self.message_id, timer=self.timer, **self.kwargs + ), + TraceTaskName.SUGGESTED_QUESTION_TRACE: lambda: self.suggested_question_trace( + message_id=self.message_id, timer=self.timer, **self.kwargs + ), + TraceTaskName.DATASET_RETRIEVAL_TRACE: lambda: self.dataset_retrieval_trace( + message_id=self.message_id, timer=self.timer, **self.kwargs + ), + TraceTaskName.TOOL_TRACE: lambda: self.tool_trace( + message_id=self.message_id, timer=self.timer, **self.kwargs + ), + TraceTaskName.GENERATE_NAME_TRACE: lambda: self.generate_name_trace( + conversation_id=self.conversation_id, timer=self.timer, **self.kwargs + ), + } + + return preprocess_map.get(self.trace_type, lambda: None)() + + # process methods for different trace types + def conversation_trace(self, **kwargs): + return kwargs + + def workflow_trace( + self, + *, + workflow_run_id: str | None, + conversation_id: str | None, + user_id: str | None, + ): + if not workflow_run_id: + return {} + + with Session(db.engine) as session: + workflow_run_stmt = select(WorkflowRun).where(WorkflowRun.id == workflow_run_id) + workflow_run = session.scalars(workflow_run_stmt).first() + if not workflow_run: + raise ValueError("Workflow run not found") + + workflow_id = workflow_run.workflow_id + tenant_id = workflow_run.tenant_id + workflow_run_id = workflow_run.id + workflow_run_elapsed_time = workflow_run.elapsed_time + workflow_run_status = workflow_run.status + workflow_run_inputs = workflow_run.inputs_dict + workflow_run_outputs = workflow_run.outputs_dict + workflow_run_version = workflow_run.version + error = workflow_run.error or "" + + total_tokens = workflow_run.total_tokens + + file_list = workflow_run_inputs.get("sys.file") or [] + query = workflow_run_inputs.get("query") or workflow_run_inputs.get("sys.query") or "" + + # get workflow_app_log_id + workflow_app_log_data_stmt = select(WorkflowAppLog.id).where( + WorkflowAppLog.tenant_id == tenant_id, + WorkflowAppLog.app_id == workflow_run.app_id, + WorkflowAppLog.workflow_run_id == workflow_run.id, + ) + workflow_app_log_id = session.scalar(workflow_app_log_data_stmt) + # get message_id + message_id = None + if conversation_id: + message_data_stmt = select(Message.id).where( + Message.conversation_id == conversation_id, + Message.workflow_run_id == workflow_run_id, + ) + message_id = session.scalar(message_data_stmt) + + metadata = { + "workflow_id": workflow_id, + "conversation_id": conversation_id, + "workflow_run_id": workflow_run_id, + "tenant_id": tenant_id, + "elapsed_time": workflow_run_elapsed_time, + "status": workflow_run_status, + "version": workflow_run_version, + "total_tokens": total_tokens, + "file_list": file_list, + "triggered_form": workflow_run.triggered_from, + "user_id": user_id, + } + + workflow_trace_info = WorkflowTraceInfo( + workflow_data=workflow_run.to_dict(), + conversation_id=conversation_id, + workflow_id=workflow_id, + tenant_id=tenant_id, + workflow_run_id=workflow_run_id, + workflow_run_elapsed_time=workflow_run_elapsed_time, + workflow_run_status=workflow_run_status, + workflow_run_inputs=workflow_run_inputs, + workflow_run_outputs=workflow_run_outputs, + workflow_run_version=workflow_run_version, + error=error, + total_tokens=total_tokens, + file_list=file_list, + query=query, + metadata=metadata, + workflow_app_log_id=workflow_app_log_id, + message_id=message_id, + start_time=workflow_run.created_at, + end_time=workflow_run.finished_at, + ) + return workflow_trace_info + + def message_trace(self, message_id: str | None): + if not message_id: + return {} + message_data = get_message_data(message_id) + if not message_data: + return {} + conversation_mode_stmt = select(Conversation.mode).where(Conversation.id == message_data.conversation_id) + conversation_mode = db.session.scalars(conversation_mode_stmt).all() + if not conversation_mode or len(conversation_mode) == 0: + return {} + conversation_mode = conversation_mode[0] + created_at = message_data.created_at + inputs = message_data.message + + # get message file data + message_file_data = db.session.query(MessageFile).filter_by(message_id=message_id).first() + file_list = [] + if message_file_data and message_file_data.url is not None: + file_url = f"{self.file_base_url}/{message_file_data.url}" if message_file_data else "" + file_list.append(file_url) + + metadata = { + "conversation_id": message_data.conversation_id, + "ls_provider": message_data.model_provider, + "ls_model_name": message_data.model_id, + "status": message_data.status, + "from_end_user_id": message_data.from_end_user_id, + "from_account_id": message_data.from_account_id, + "agent_based": message_data.agent_based, + "workflow_run_id": message_data.workflow_run_id, + "from_source": message_data.from_source, + "message_id": message_id, + } + + message_tokens = message_data.message_tokens + + message_trace_info = MessageTraceInfo( + message_id=message_id, + message_data=message_data.to_dict(), + conversation_model=conversation_mode, + message_tokens=message_tokens, + answer_tokens=message_data.answer_tokens, + total_tokens=message_tokens + message_data.answer_tokens, + error=message_data.error or "", + inputs=inputs, + outputs=message_data.answer, + file_list=file_list, + start_time=created_at, + end_time=created_at + timedelta(seconds=message_data.provider_response_latency), + metadata=metadata, + message_file_data=message_file_data, + conversation_mode=conversation_mode, + ) + + return message_trace_info + + def moderation_trace(self, message_id, timer, **kwargs): + moderation_result = kwargs.get("moderation_result") + if not moderation_result: + return {} + inputs = kwargs.get("inputs") + message_data = get_message_data(message_id) + if not message_data: + return {} + metadata = { + "message_id": message_id, + "action": moderation_result.action, + "preset_response": moderation_result.preset_response, + "query": moderation_result.query, + } + + # get workflow_app_log_id + workflow_app_log_id = None + if message_data.workflow_run_id: + workflow_app_log_data = ( + db.session.query(WorkflowAppLog).filter_by(workflow_run_id=message_data.workflow_run_id).first() + ) + workflow_app_log_id = str(workflow_app_log_data.id) if workflow_app_log_data else None + + moderation_trace_info = ModerationTraceInfo( + message_id=workflow_app_log_id or message_id, + inputs=inputs, + message_data=message_data.to_dict(), + flagged=moderation_result.flagged, + action=moderation_result.action, + preset_response=moderation_result.preset_response, + query=moderation_result.query, + start_time=timer.get("start"), + end_time=timer.get("end"), + metadata=metadata, + ) + + return moderation_trace_info + + def suggested_question_trace(self, message_id, timer, **kwargs): + suggested_question = kwargs.get("suggested_question", []) + message_data = get_message_data(message_id) + if not message_data: + return {} + metadata = { + "message_id": message_id, + "ls_provider": message_data.model_provider, + "ls_model_name": message_data.model_id, + "status": message_data.status, + "from_end_user_id": message_data.from_end_user_id, + "from_account_id": message_data.from_account_id, + "agent_based": message_data.agent_based, + "workflow_run_id": message_data.workflow_run_id, + "from_source": message_data.from_source, + } + + # get workflow_app_log_id + workflow_app_log_id = None + if message_data.workflow_run_id: + workflow_app_log_data = ( + db.session.query(WorkflowAppLog).filter_by(workflow_run_id=message_data.workflow_run_id).first() + ) + workflow_app_log_id = str(workflow_app_log_data.id) if workflow_app_log_data else None + + suggested_question_trace_info = SuggestedQuestionTraceInfo( + message_id=workflow_app_log_id or message_id, + message_data=message_data.to_dict(), + inputs=message_data.message, + outputs=message_data.answer, + start_time=timer.get("start"), + end_time=timer.get("end"), + metadata=metadata, + total_tokens=message_data.message_tokens + message_data.answer_tokens, + status=message_data.status, + error=message_data.error, + from_account_id=message_data.from_account_id, + agent_based=message_data.agent_based, + from_source=message_data.from_source, + model_provider=message_data.model_provider, + model_id=message_data.model_id, + suggested_question=suggested_question, + level=message_data.status, + status_message=message_data.error, + ) + + return suggested_question_trace_info + + def dataset_retrieval_trace(self, message_id, timer, **kwargs): + documents = kwargs.get("documents") + message_data = get_message_data(message_id) + if not message_data: + return {} + + metadata = { + "message_id": message_id, + "ls_provider": message_data.model_provider, + "ls_model_name": message_data.model_id, + "status": message_data.status, + "from_end_user_id": message_data.from_end_user_id, + "from_account_id": message_data.from_account_id, + "agent_based": message_data.agent_based, + "workflow_run_id": message_data.workflow_run_id, + "from_source": message_data.from_source, + } + + dataset_retrieval_trace_info = DatasetRetrievalTraceInfo( + message_id=message_id, + inputs=message_data.query or message_data.inputs, + documents=[doc.model_dump() for doc in documents] if documents else [], + start_time=timer.get("start"), + end_time=timer.get("end"), + metadata=metadata, + message_data=message_data.to_dict(), + ) + + return dataset_retrieval_trace_info + + def tool_trace(self, message_id, timer, **kwargs): + tool_name = kwargs.get("tool_name", "") + tool_inputs = kwargs.get("tool_inputs", {}) + tool_outputs = kwargs.get("tool_outputs", {}) + message_data = get_message_data(message_id) + if not message_data: + return {} + tool_config = {} + time_cost = 0 + error = None + tool_parameters = {} + created_time = message_data.created_at + end_time = message_data.updated_at + agent_thoughts = message_data.agent_thoughts + for agent_thought in agent_thoughts: + if tool_name in agent_thought.tools: + created_time = agent_thought.created_at + tool_meta_data = agent_thought.tool_meta.get(tool_name, {}) + tool_config = tool_meta_data.get("tool_config", {}) + time_cost = tool_meta_data.get("time_cost", 0) + end_time = created_time + timedelta(seconds=time_cost) + error = tool_meta_data.get("error", "") + tool_parameters = tool_meta_data.get("tool_parameters", {}) + metadata = { + "message_id": message_id, + "tool_name": tool_name, + "tool_inputs": tool_inputs, + "tool_outputs": tool_outputs, + "tool_config": tool_config, + "time_cost": time_cost, + "error": error, + "tool_parameters": tool_parameters, + } + + file_url = "" + message_file_data = db.session.query(MessageFile).filter_by(message_id=message_id).first() + if message_file_data: + message_file_id = message_file_data.id if message_file_data else None + type = message_file_data.type + created_by_role = message_file_data.created_by_role + created_user_id = message_file_data.created_by + file_url = f"{self.file_base_url}/{message_file_data.url}" + + metadata.update( + { + "message_file_id": message_file_id, + "created_by_role": created_by_role, + "created_user_id": created_user_id, + "type": type, + } + ) + + tool_trace_info = ToolTraceInfo( + message_id=message_id, + message_data=message_data.to_dict(), + tool_name=tool_name, + start_time=timer.get("start") if timer else created_time, + end_time=timer.get("end") if timer else end_time, + tool_inputs=tool_inputs, + tool_outputs=tool_outputs, + metadata=metadata, + message_file_data=message_file_data, + error=error, + inputs=message_data.message, + outputs=message_data.answer, + tool_config=tool_config, + time_cost=time_cost, + tool_parameters=tool_parameters, + file_url=file_url, + ) + + return tool_trace_info + + def generate_name_trace(self, conversation_id, timer, **kwargs): + generate_conversation_name = kwargs.get("generate_conversation_name") + inputs = kwargs.get("inputs") + tenant_id = kwargs.get("tenant_id") + if not tenant_id: + return {} + start_time = timer.get("start") + end_time = timer.get("end") + + metadata = { + "conversation_id": conversation_id, + "tenant_id": tenant_id, + } + + generate_name_trace_info = GenerateNameTraceInfo( + conversation_id=conversation_id, + inputs=inputs, + outputs=generate_conversation_name, + start_time=start_time, + end_time=end_time, + metadata=metadata, + tenant_id=tenant_id, + ) + + return generate_name_trace_info + + +trace_manager_timer: Optional[threading.Timer] = None +trace_manager_queue: queue.Queue = queue.Queue() +trace_manager_interval = int(os.getenv("TRACE_QUEUE_MANAGER_INTERVAL", 5)) +trace_manager_batch_size = int(os.getenv("TRACE_QUEUE_MANAGER_BATCH_SIZE", 100)) + + +class TraceQueueManager: + def __init__(self, app_id=None, user_id=None): + global trace_manager_timer + + self.app_id = app_id + self.user_id = user_id + self.trace_instance = OpsTraceManager.get_ops_trace_instance(app_id) + self.flask_app = current_app._get_current_object() # type: ignore + if trace_manager_timer is None: + self.start_timer() + + def add_trace_task(self, trace_task: TraceTask): + global trace_manager_timer, trace_manager_queue + try: + if self.trace_instance: + trace_task.app_id = self.app_id + trace_manager_queue.put(trace_task) + except Exception as e: + logging.exception(f"Error adding trace task, trace_type {trace_task.trace_type}") + finally: + self.start_timer() + + def collect_tasks(self): + global trace_manager_queue + tasks: list[TraceTask] = [] + while len(tasks) < trace_manager_batch_size and not trace_manager_queue.empty(): + task = trace_manager_queue.get_nowait() + tasks.append(task) + trace_manager_queue.task_done() + return tasks + + def run(self): + try: + tasks = self.collect_tasks() + if tasks: + self.send_to_celery(tasks) + except Exception as e: + logging.exception("Error processing trace tasks") + + def start_timer(self): + global trace_manager_timer + if trace_manager_timer is None or not trace_manager_timer.is_alive(): + trace_manager_timer = threading.Timer(trace_manager_interval, self.run) + trace_manager_timer.name = f"trace_manager_timer_{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}" + trace_manager_timer.daemon = False + trace_manager_timer.start() + + def send_to_celery(self, tasks: list[TraceTask]): + with self.flask_app.app_context(): + for task in tasks: + if task.app_id is None: + continue + file_id = uuid4().hex + trace_info = task.execute() + task_data = TaskData( + app_id=task.app_id, + trace_info_type=type(trace_info).__name__, + trace_info=trace_info.model_dump() if trace_info else None, + ) + file_path = f"{OPS_FILE_PATH}{task.app_id}/{file_id}.json" + storage.save(file_path, task_data.model_dump_json().encode("utf-8")) + file_info = { + "file_id": file_id, + "app_id": task.app_id, + } + process_trace_tasks.delay(file_info) diff --git a/api/core/ops/utils.py b/api/core/ops/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..8b06df1930595d8e30bb285e76967b8b3a057bfa --- /dev/null +++ b/api/core/ops/utils.py @@ -0,0 +1,62 @@ +from contextlib import contextmanager +from datetime import datetime +from typing import Optional, Union + +from extensions.ext_database import db +from models.model import Message + + +def filter_none_values(data: dict): + new_data = {} + for key, value in data.items(): + if value is None: + continue + if isinstance(value, datetime): + new_data[key] = value.isoformat() + else: + new_data[key] = value + return new_data + + +def get_message_data(message_id: str): + return db.session.query(Message).filter(Message.id == message_id).first() + + +@contextmanager +def measure_time(): + timing_info = {"start": datetime.now(), "end": None} + try: + yield timing_info + finally: + timing_info["end"] = datetime.now() + + +def replace_text_with_content(data): + if isinstance(data, dict): + new_data = {} + for key, value in data.items(): + if key == "text": + new_data["content"] = value + else: + new_data[key] = replace_text_with_content(value) + return new_data + elif isinstance(data, list): + return [replace_text_with_content(item) for item in data] + else: + return data + + +def generate_dotted_order( + run_id: str, start_time: Union[str, datetime], parent_dotted_order: Optional[str] = None +) -> str: + """ + generate dotted_order for langsmith + """ + start_time = datetime.fromisoformat(start_time) if isinstance(start_time, str) else start_time + timestamp = start_time.strftime("%Y%m%dT%H%M%S%f")[:-3] + "Z" + current_segment = f"{timestamp}{run_id}" + + if parent_dotted_order is None: + return current_segment + + return f"{parent_dotted_order}.{current_segment}" diff --git a/api/core/prompt/__init__.py b/api/core/prompt/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/prompt/advanced_prompt_transform.py b/api/core/prompt/advanced_prompt_transform.py new file mode 100644 index 0000000000000000000000000000000000000000..87c7a79fb0120143847e742b35a84cc98a44ccf3 --- /dev/null +++ b/api/core/prompt/advanced_prompt_transform.py @@ -0,0 +1,287 @@ +from collections.abc import Mapping, Sequence +from typing import Optional, cast + +from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity +from core.file import file_manager +from core.file.models import File +from core.helper.code_executor.jinja2.jinja2_formatter import Jinja2Formatter +from core.memory.token_buffer_memory import TokenBufferMemory +from core.model_runtime.entities import ( + AssistantPromptMessage, + PromptMessage, + PromptMessageContent, + PromptMessageRole, + SystemPromptMessage, + TextPromptMessageContent, + UserPromptMessage, +) +from core.model_runtime.entities.message_entities import ImagePromptMessageContent +from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate, MemoryConfig +from core.prompt.prompt_transform import PromptTransform +from core.prompt.utils.prompt_template_parser import PromptTemplateParser +from core.workflow.entities.variable_pool import VariablePool + + +class AdvancedPromptTransform(PromptTransform): + """ + Advanced Prompt Transform for Workflow LLM Node. + """ + + def __init__( + self, + with_variable_tmpl: bool = False, + image_detail_config: ImagePromptMessageContent.DETAIL = ImagePromptMessageContent.DETAIL.LOW, + ) -> None: + self.with_variable_tmpl = with_variable_tmpl + self.image_detail_config = image_detail_config + + def get_prompt( + self, + *, + prompt_template: Sequence[ChatModelMessage] | CompletionModelPromptTemplate, + inputs: Mapping[str, str], + query: str, + files: Sequence[File], + context: Optional[str], + memory_config: Optional[MemoryConfig], + memory: Optional[TokenBufferMemory], + model_config: ModelConfigWithCredentialsEntity, + ) -> list[PromptMessage]: + prompt_messages = [] + + if isinstance(prompt_template, CompletionModelPromptTemplate): + prompt_messages = self._get_completion_model_prompt_messages( + prompt_template=prompt_template, + inputs=inputs, + query=query, + files=files, + context=context, + memory_config=memory_config, + memory=memory, + model_config=model_config, + ) + elif isinstance(prompt_template, list) and all(isinstance(item, ChatModelMessage) for item in prompt_template): + prompt_messages = self._get_chat_model_prompt_messages( + prompt_template=prompt_template, + inputs=inputs, + query=query, + files=files, + context=context, + memory_config=memory_config, + memory=memory, + model_config=model_config, + ) + + return prompt_messages + + def _get_completion_model_prompt_messages( + self, + prompt_template: CompletionModelPromptTemplate, + inputs: Mapping[str, str], + query: Optional[str], + files: Sequence[File], + context: Optional[str], + memory_config: Optional[MemoryConfig], + memory: Optional[TokenBufferMemory], + model_config: ModelConfigWithCredentialsEntity, + ) -> list[PromptMessage]: + """ + Get completion model prompt messages. + """ + raw_prompt = prompt_template.text + + prompt_messages: list[PromptMessage] = [] + + if prompt_template.edition_type == "basic" or not prompt_template.edition_type: + parser = PromptTemplateParser(template=raw_prompt, with_variable_tmpl=self.with_variable_tmpl) + prompt_inputs: Mapping[str, str] = {k: inputs[k] for k in parser.variable_keys if k in inputs} + + prompt_inputs = self._set_context_variable(context, parser, prompt_inputs) + + if memory and memory_config and memory_config.role_prefix: + role_prefix = memory_config.role_prefix + prompt_inputs = self._set_histories_variable( + memory=memory, + memory_config=memory_config, + raw_prompt=raw_prompt, + role_prefix=role_prefix, + parser=parser, + prompt_inputs=prompt_inputs, + model_config=model_config, + ) + + if query: + prompt_inputs = self._set_query_variable(query, parser, prompt_inputs) + + prompt = parser.format(prompt_inputs) + else: + prompt = raw_prompt + prompt_inputs = inputs + + prompt = Jinja2Formatter.format(prompt, prompt_inputs) + + if files: + prompt_message_contents: list[PromptMessageContent] = [] + prompt_message_contents.append(TextPromptMessageContent(data=prompt)) + for file in files: + prompt_message_contents.append(file_manager.to_prompt_message_content(file)) + + prompt_messages.append(UserPromptMessage(content=prompt_message_contents)) + else: + prompt_messages.append(UserPromptMessage(content=prompt)) + + return prompt_messages + + def _get_chat_model_prompt_messages( + self, + prompt_template: list[ChatModelMessage], + inputs: Mapping[str, str], + query: Optional[str], + files: Sequence[File], + context: Optional[str], + memory_config: Optional[MemoryConfig], + memory: Optional[TokenBufferMemory], + model_config: ModelConfigWithCredentialsEntity, + ) -> list[PromptMessage]: + """ + Get chat model prompt messages. + """ + prompt_messages: list[PromptMessage] = [] + for prompt_item in prompt_template: + raw_prompt = prompt_item.text + + if prompt_item.edition_type == "basic" or not prompt_item.edition_type: + if self.with_variable_tmpl: + vp = VariablePool() + for k, v in inputs.items(): + if k.startswith("#"): + vp.add(k[1:-1].split("."), v) + raw_prompt = raw_prompt.replace("{{#context#}}", context or "") + prompt = vp.convert_template(raw_prompt).text + else: + parser = PromptTemplateParser(template=raw_prompt, with_variable_tmpl=self.with_variable_tmpl) + prompt_inputs: Mapping[str, str] = {k: inputs[k] for k in parser.variable_keys if k in inputs} + prompt_inputs = self._set_context_variable( + context=context, parser=parser, prompt_inputs=prompt_inputs + ) + prompt = parser.format(prompt_inputs) + elif prompt_item.edition_type == "jinja2": + prompt = raw_prompt + prompt_inputs = inputs + prompt = Jinja2Formatter.format(template=prompt, inputs=prompt_inputs) + else: + raise ValueError(f"Invalid edition type: {prompt_item.edition_type}") + + if prompt_item.role == PromptMessageRole.USER: + prompt_messages.append(UserPromptMessage(content=prompt)) + elif prompt_item.role == PromptMessageRole.SYSTEM and prompt: + prompt_messages.append(SystemPromptMessage(content=prompt)) + elif prompt_item.role == PromptMessageRole.ASSISTANT: + prompt_messages.append(AssistantPromptMessage(content=prompt)) + + if query and memory_config and memory_config.query_prompt_template: + parser = PromptTemplateParser( + template=memory_config.query_prompt_template, with_variable_tmpl=self.with_variable_tmpl + ) + prompt_inputs = {k: inputs[k] for k in parser.variable_keys if k in inputs} + prompt_inputs["#sys.query#"] = query + + prompt_inputs = self._set_context_variable(context, parser, prompt_inputs) + + query = parser.format(prompt_inputs) + + if memory and memory_config: + prompt_messages = self._append_chat_histories(memory, memory_config, prompt_messages, model_config) + + if files and query is not None: + prompt_message_contents: list[PromptMessageContent] = [] + prompt_message_contents.append(TextPromptMessageContent(data=query)) + for file in files: + prompt_message_contents.append(file_manager.to_prompt_message_content(file)) + prompt_messages.append(UserPromptMessage(content=prompt_message_contents)) + else: + prompt_messages.append(UserPromptMessage(content=query)) + elif files: + if not query: + # get last message + last_message = prompt_messages[-1] if prompt_messages else None + if last_message and last_message.role == PromptMessageRole.USER: + # get last user message content and add files + prompt_message_contents = [TextPromptMessageContent(data=cast(str, last_message.content))] + for file in files: + prompt_message_contents.append(file_manager.to_prompt_message_content(file)) + + last_message.content = prompt_message_contents + else: + prompt_message_contents = [TextPromptMessageContent(data="")] # not for query + for file in files: + prompt_message_contents.append(file_manager.to_prompt_message_content(file)) + + prompt_messages.append(UserPromptMessage(content=prompt_message_contents)) + else: + prompt_message_contents = [TextPromptMessageContent(data=query)] + for file in files: + prompt_message_contents.append(file_manager.to_prompt_message_content(file)) + + prompt_messages.append(UserPromptMessage(content=prompt_message_contents)) + elif query: + prompt_messages.append(UserPromptMessage(content=query)) + + return prompt_messages + + def _set_context_variable( + self, context: str | None, parser: PromptTemplateParser, prompt_inputs: Mapping[str, str] + ) -> Mapping[str, str]: + prompt_inputs = dict(prompt_inputs) + if "#context#" in parser.variable_keys: + if context: + prompt_inputs["#context#"] = context + else: + prompt_inputs["#context#"] = "" + + return prompt_inputs + + def _set_query_variable( + self, query: str, parser: PromptTemplateParser, prompt_inputs: Mapping[str, str] + ) -> Mapping[str, str]: + prompt_inputs = dict(prompt_inputs) + if "#query#" in parser.variable_keys: + if query: + prompt_inputs["#query#"] = query + else: + prompt_inputs["#query#"] = "" + + return prompt_inputs + + def _set_histories_variable( + self, + memory: TokenBufferMemory, + memory_config: MemoryConfig, + raw_prompt: str, + role_prefix: MemoryConfig.RolePrefix, + parser: PromptTemplateParser, + prompt_inputs: Mapping[str, str], + model_config: ModelConfigWithCredentialsEntity, + ) -> Mapping[str, str]: + prompt_inputs = dict(prompt_inputs) + if "#histories#" in parser.variable_keys: + if memory: + inputs = {"#histories#": "", **prompt_inputs} + parser = PromptTemplateParser(template=raw_prompt, with_variable_tmpl=self.with_variable_tmpl) + prompt_inputs = {k: inputs[k] for k in parser.variable_keys if k in inputs} + tmp_human_message = UserPromptMessage(content=parser.format(prompt_inputs)) + + rest_tokens = self._calculate_rest_token([tmp_human_message], model_config) + + histories = self._get_history_messages_from_memory( + memory=memory, + memory_config=memory_config, + max_token_limit=rest_tokens, + human_prefix=role_prefix.user, + ai_prefix=role_prefix.assistant, + ) + prompt_inputs["#histories#"] = histories + else: + prompt_inputs["#histories#"] = "" + + return prompt_inputs diff --git a/api/core/prompt/agent_history_prompt_transform.py b/api/core/prompt/agent_history_prompt_transform.py new file mode 100644 index 0000000000000000000000000000000000000000..09f017a7db0d3aa661b84057bd82f5495f3a72c1 --- /dev/null +++ b/api/core/prompt/agent_history_prompt_transform.py @@ -0,0 +1,80 @@ +from typing import Optional, cast + +from core.app.entities.app_invoke_entities import ( + ModelConfigWithCredentialsEntity, +) +from core.memory.token_buffer_memory import TokenBufferMemory +from core.model_runtime.entities.message_entities import ( + PromptMessage, + SystemPromptMessage, + UserPromptMessage, +) +from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel +from core.prompt.prompt_transform import PromptTransform + + +class AgentHistoryPromptTransform(PromptTransform): + """ + History Prompt Transform for Agent App + """ + + def __init__( + self, + model_config: ModelConfigWithCredentialsEntity, + prompt_messages: list[PromptMessage], + history_messages: list[PromptMessage], + memory: Optional[TokenBufferMemory] = None, + ): + self.model_config = model_config + self.prompt_messages = prompt_messages + self.history_messages = history_messages + self.memory = memory + + def get_prompt(self) -> list[PromptMessage]: + prompt_messages: list[PromptMessage] = [] + num_system = 0 + for prompt_message in self.history_messages: + if isinstance(prompt_message, SystemPromptMessage): + prompt_messages.append(prompt_message) + num_system += 1 + + if not self.memory: + return prompt_messages + + max_token_limit = self._calculate_rest_token(self.prompt_messages, self.model_config) + + model_type_instance = self.model_config.provider_model_bundle.model_type_instance + model_type_instance = cast(LargeLanguageModel, model_type_instance) + + curr_message_tokens = model_type_instance.get_num_tokens( + self.memory.model_instance.model, self.memory.model_instance.credentials, self.history_messages + ) + if curr_message_tokens <= max_token_limit: + return self.history_messages + + # number of prompt has been appended in current message + num_prompt = 0 + # append prompt messages in desc order + for prompt_message in self.history_messages[::-1]: + if isinstance(prompt_message, SystemPromptMessage): + continue + prompt_messages.append(prompt_message) + num_prompt += 1 + # a message is start with UserPromptMessage + if isinstance(prompt_message, UserPromptMessage): + curr_message_tokens = model_type_instance.get_num_tokens( + self.memory.model_instance.model, self.memory.model_instance.credentials, prompt_messages + ) + # if current message token is overflow, drop all the prompts in current message and break + if curr_message_tokens > max_token_limit: + prompt_messages = prompt_messages[:-num_prompt] + break + num_prompt = 0 + # return prompt messages in asc order + message_prompts = prompt_messages[num_system:] + message_prompts.reverse() + + # merge system and message prompt + prompt_messages = prompt_messages[:num_system] + prompt_messages.extend(message_prompts) + return prompt_messages diff --git a/api/core/prompt/entities/__init__.py b/api/core/prompt/entities/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/prompt/entities/advanced_prompt_entities.py b/api/core/prompt/entities/advanced_prompt_entities.py new file mode 100644 index 0000000000000000000000000000000000000000..c8e7b414dffe8550ed2c078c4c779641e043cfd6 --- /dev/null +++ b/api/core/prompt/entities/advanced_prompt_entities.py @@ -0,0 +1,50 @@ +from typing import Literal, Optional + +from pydantic import BaseModel + +from core.model_runtime.entities.message_entities import PromptMessageRole + + +class ChatModelMessage(BaseModel): + """ + Chat Message. + """ + + text: str + role: PromptMessageRole + edition_type: Optional[Literal["basic", "jinja2"]] = None + + +class CompletionModelPromptTemplate(BaseModel): + """ + Completion Model Prompt Template. + """ + + text: str + edition_type: Optional[Literal["basic", "jinja2"]] = None + + +class MemoryConfig(BaseModel): + """ + Memory Config. + """ + + class RolePrefix(BaseModel): + """ + Role Prefix. + """ + + user: str + assistant: str + + class WindowConfig(BaseModel): + """ + Window Config. + """ + + enabled: bool + size: Optional[int] = None + + role_prefix: Optional[RolePrefix] = None + window: WindowConfig + query_prompt_template: Optional[str] = None diff --git a/api/core/prompt/prompt_templates/__init__.py b/api/core/prompt/prompt_templates/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/prompt/prompt_templates/advanced_prompt_templates.py b/api/core/prompt/prompt_templates/advanced_prompt_templates.py new file mode 100644 index 0000000000000000000000000000000000000000..e55966eeee8bbda89e4d212f8c3d1a2734fe9b8c --- /dev/null +++ b/api/core/prompt/prompt_templates/advanced_prompt_templates.py @@ -0,0 +1,45 @@ +CONTEXT = "Use the following context as your learned knowledge, inside XML tags.\n\n\n{{#context#}}\n\n\nWhen answer to user:\n- If you don't know, just say that you don't know.\n- If you don't know when you are not sure, ask for clarification.\nAvoid mentioning that you obtained the information from the context.\nAnd answer according to the language of the user's question.\n" # noqa: E501 + +BAICHUAN_CONTEXT = "用户在与一个客观的助手对话。助手会尊重找到的材料,给出全面专业的解释,但不会过度演绎。同时回答中不会暴露引用的材料:\n\n```\n{{#context#}}\n```\n\n" # noqa: E501 + +CHAT_APP_COMPLETION_PROMPT_CONFIG = { + "completion_prompt_config": { + "prompt": { + "text": "{{#pre_prompt#}}\nHere are the chat histories between human and assistant, inside XML tags.\n\n\n{{#histories#}}\n\n\n\nHuman: {{#query#}}\n\nAssistant: " # noqa: E501 + }, + "conversation_histories_role": {"user_prefix": "Human", "assistant_prefix": "Assistant"}, + }, + "stop": ["Human:"], +} + +CHAT_APP_CHAT_PROMPT_CONFIG = {"chat_prompt_config": {"prompt": [{"role": "system", "text": "{{#pre_prompt#}}"}]}} + +COMPLETION_APP_CHAT_PROMPT_CONFIG = {"chat_prompt_config": {"prompt": [{"role": "user", "text": "{{#pre_prompt#}}"}]}} + +COMPLETION_APP_COMPLETION_PROMPT_CONFIG = { + "completion_prompt_config": {"prompt": {"text": "{{#pre_prompt#}}"}}, + "stop": ["Human:"], +} + +BAICHUAN_CHAT_APP_COMPLETION_PROMPT_CONFIG = { + "completion_prompt_config": { + "prompt": { + "text": "{{#pre_prompt#}}\n\n用户和助手的历史对话内容如下:\n```\n{{#histories#}}\n```\n\n\n\n用户:{{#query#}}" # noqa: E501 + }, + "conversation_histories_role": {"user_prefix": "用户", "assistant_prefix": "助手"}, + }, + "stop": ["用户:"], +} + +BAICHUAN_CHAT_APP_CHAT_PROMPT_CONFIG = { + "chat_prompt_config": {"prompt": [{"role": "system", "text": "{{#pre_prompt#}}"}]} +} + +BAICHUAN_COMPLETION_APP_CHAT_PROMPT_CONFIG = { + "chat_prompt_config": {"prompt": [{"role": "user", "text": "{{#pre_prompt#}}"}]} +} + +BAICHUAN_COMPLETION_APP_COMPLETION_PROMPT_CONFIG = { + "completion_prompt_config": {"prompt": {"text": "{{#pre_prompt#}}"}}, + "stop": ["用户:"], +} diff --git a/api/core/prompt/prompt_templates/baichuan_chat.json b/api/core/prompt/prompt_templates/baichuan_chat.json new file mode 100644 index 0000000000000000000000000000000000000000..03b6a53cfff2d1fda9e732b022988d0a886175c3 --- /dev/null +++ b/api/core/prompt/prompt_templates/baichuan_chat.json @@ -0,0 +1,13 @@ +{ + "human_prefix": "用户", + "assistant_prefix": "助手", + "context_prompt": "用户在与一个客观的助手对话。助手会尊重找到的材料,给出全面专业的解释,但不会过度演绎。同时回答中不会暴露引用的材料:\n\n```\n{{#context#}}\n```\n\n", + "histories_prompt": "用户和助手的历史对话内容如下:\n```\n{{#histories#}}\n```\n\n", + "system_prompt_orders": [ + "context_prompt", + "pre_prompt", + "histories_prompt" + ], + "query_prompt": "\n\n用户:{{#query#}}", + "stops": ["用户:"] +} \ No newline at end of file diff --git a/api/core/prompt/prompt_templates/baichuan_completion.json b/api/core/prompt/prompt_templates/baichuan_completion.json new file mode 100644 index 0000000000000000000000000000000000000000..ae8c0dac53392fe50de3845343659a6fe8de98ae --- /dev/null +++ b/api/core/prompt/prompt_templates/baichuan_completion.json @@ -0,0 +1,9 @@ +{ + "context_prompt": "用户在与一个客观的助手对话。助手会尊重找到的材料,给出全面专业的解释,但不会过度演绎。同时回答中不会暴露引用的材料:\n\n```\n{{#context#}}\n```\n", + "system_prompt_orders": [ + "context_prompt", + "pre_prompt" + ], + "query_prompt": "{{#query#}}", + "stops": null +} \ No newline at end of file diff --git a/api/core/prompt/prompt_templates/common_chat.json b/api/core/prompt/prompt_templates/common_chat.json new file mode 100644 index 0000000000000000000000000000000000000000..d398a512e670a72f18e3a56b18af8947ae6fd2ba --- /dev/null +++ b/api/core/prompt/prompt_templates/common_chat.json @@ -0,0 +1,13 @@ +{ + "human_prefix": "Human", + "assistant_prefix": "Assistant", + "context_prompt": "Use the following context as your learned knowledge, inside XML tags.\n\n\n{{#context#}}\n\n\nWhen answer to user:\n- If you don't know, just say that you don't know.\n- If you don't know when you are not sure, ask for clarification.\nAvoid mentioning that you obtained the information from the context.\nAnd answer according to the language of the user's question.\n\n", + "histories_prompt": "Here is the chat histories between human and assistant, inside XML tags.\n\n\n{{#histories#}}\n\n\n", + "system_prompt_orders": [ + "context_prompt", + "pre_prompt", + "histories_prompt" + ], + "query_prompt": "\n\nHuman: {{#query#}}\n\nAssistant: ", + "stops": ["\nHuman:", ""] +} diff --git a/api/core/prompt/prompt_templates/common_completion.json b/api/core/prompt/prompt_templates/common_completion.json new file mode 100644 index 0000000000000000000000000000000000000000..c148772010fb059d029357124640416ea6dd7a07 --- /dev/null +++ b/api/core/prompt/prompt_templates/common_completion.json @@ -0,0 +1,9 @@ +{ + "context_prompt": "Use the following context as your learned knowledge, inside XML tags.\n\n\n{{#context#}}\n\n\nWhen answer to user:\n- If you don't know, just say that you don't know.\n- If you don't know when you are not sure, ask for clarification.\nAvoid mentioning that you obtained the information from the context.\nAnd answer according to the language of the user's question.\n\n", + "system_prompt_orders": [ + "context_prompt", + "pre_prompt" + ], + "query_prompt": "{{#query#}}", + "stops": null +} \ No newline at end of file diff --git a/api/core/prompt/prompt_transform.py b/api/core/prompt/prompt_transform.py new file mode 100644 index 0000000000000000000000000000000000000000..1f040599be6dac390a5f2250ba853740b57a76c5 --- /dev/null +++ b/api/core/prompt/prompt_transform.py @@ -0,0 +1,90 @@ +from typing import Any, Optional + +from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity +from core.memory.token_buffer_memory import TokenBufferMemory +from core.model_manager import ModelInstance +from core.model_runtime.entities.message_entities import PromptMessage +from core.model_runtime.entities.model_entities import ModelPropertyKey +from core.prompt.entities.advanced_prompt_entities import MemoryConfig + + +class PromptTransform: + def _append_chat_histories( + self, + memory: TokenBufferMemory, + memory_config: MemoryConfig, + prompt_messages: list[PromptMessage], + model_config: ModelConfigWithCredentialsEntity, + ) -> list[PromptMessage]: + rest_tokens = self._calculate_rest_token(prompt_messages, model_config) + histories = self._get_history_messages_list_from_memory(memory, memory_config, rest_tokens) + prompt_messages.extend(histories) + + return prompt_messages + + def _calculate_rest_token( + self, prompt_messages: list[PromptMessage], model_config: ModelConfigWithCredentialsEntity + ) -> int: + rest_tokens = 2000 + + model_context_tokens = model_config.model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE) + if model_context_tokens: + model_instance = ModelInstance( + provider_model_bundle=model_config.provider_model_bundle, model=model_config.model + ) + + curr_message_tokens = model_instance.get_llm_num_tokens(prompt_messages) + + max_tokens = 0 + for parameter_rule in model_config.model_schema.parameter_rules: + if parameter_rule.name == "max_tokens" or ( + parameter_rule.use_template and parameter_rule.use_template == "max_tokens" + ): + max_tokens = ( + model_config.parameters.get(parameter_rule.name) + or model_config.parameters.get(parameter_rule.use_template or "") + ) or 0 + + rest_tokens = model_context_tokens - max_tokens - curr_message_tokens + rest_tokens = max(rest_tokens, 0) + + return rest_tokens + + def _get_history_messages_from_memory( + self, + memory: TokenBufferMemory, + memory_config: MemoryConfig, + max_token_limit: int, + human_prefix: Optional[str] = None, + ai_prefix: Optional[str] = None, + ) -> str: + """Get memory messages.""" + kwargs: dict[str, Any] = {"max_token_limit": max_token_limit} + + if human_prefix: + kwargs["human_prefix"] = human_prefix + + if ai_prefix: + kwargs["ai_prefix"] = ai_prefix + + if memory_config.window.enabled and memory_config.window.size is not None and memory_config.window.size > 0: + kwargs["message_limit"] = memory_config.window.size + + return memory.get_history_prompt_text(**kwargs) + + def _get_history_messages_list_from_memory( + self, memory: TokenBufferMemory, memory_config: MemoryConfig, max_token_limit: int + ) -> list[PromptMessage]: + """Get memory messages.""" + return list( + memory.get_history_prompt_messages( + max_token_limit=max_token_limit, + message_limit=memory_config.window.size + if ( + memory_config.window.enabled + and memory_config.window.size is not None + and memory_config.window.size > 0 + ) + else None, + ) + ) diff --git a/api/core/prompt/simple_prompt_transform.py b/api/core/prompt/simple_prompt_transform.py new file mode 100644 index 0000000000000000000000000000000000000000..e75877de9b695c50a6e3f519d4c0dbb5a2091571 --- /dev/null +++ b/api/core/prompt/simple_prompt_transform.py @@ -0,0 +1,327 @@ +import enum +import json +import os +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, Optional, cast + +from core.app.app_config.entities import PromptTemplateEntity +from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity +from core.file import file_manager +from core.memory.token_buffer_memory import TokenBufferMemory +from core.model_runtime.entities.message_entities import ( + PromptMessage, + PromptMessageContent, + SystemPromptMessage, + TextPromptMessageContent, + UserPromptMessage, +) +from core.prompt.entities.advanced_prompt_entities import MemoryConfig +from core.prompt.prompt_transform import PromptTransform +from core.prompt.utils.prompt_template_parser import PromptTemplateParser +from models.model import AppMode + +if TYPE_CHECKING: + from core.file.models import File + + +class ModelMode(enum.StrEnum): + COMPLETION = "completion" + CHAT = "chat" + + @classmethod + def value_of(cls, value: str) -> "ModelMode": + """ + Get value of given mode. + + :param value: mode value + :return: mode + """ + for mode in cls: + if mode.value == value: + return mode + raise ValueError(f"invalid mode value {value}") + + +prompt_file_contents: dict[str, Any] = {} + + +class SimplePromptTransform(PromptTransform): + """ + Simple Prompt Transform for Chatbot App Basic Mode. + """ + + def get_prompt( + self, + app_mode: AppMode, + prompt_template_entity: PromptTemplateEntity, + inputs: Mapping[str, str], + query: str, + files: Sequence["File"], + context: Optional[str], + memory: Optional[TokenBufferMemory], + model_config: ModelConfigWithCredentialsEntity, + ) -> tuple[list[PromptMessage], Optional[list[str]]]: + inputs = {key: str(value) for key, value in inputs.items()} + + model_mode = ModelMode.value_of(model_config.mode) + if model_mode == ModelMode.CHAT: + prompt_messages, stops = self._get_chat_model_prompt_messages( + app_mode=app_mode, + pre_prompt=prompt_template_entity.simple_prompt_template or "", + inputs=inputs, + query=query, + files=files, + context=context, + memory=memory, + model_config=model_config, + ) + else: + prompt_messages, stops = self._get_completion_model_prompt_messages( + app_mode=app_mode, + pre_prompt=prompt_template_entity.simple_prompt_template or "", + inputs=inputs, + query=query, + files=files, + context=context, + memory=memory, + model_config=model_config, + ) + + return prompt_messages, stops + + def get_prompt_str_and_rules( + self, + app_mode: AppMode, + model_config: ModelConfigWithCredentialsEntity, + pre_prompt: str, + inputs: dict, + query: Optional[str] = None, + context: Optional[str] = None, + histories: Optional[str] = None, + ) -> tuple[str, dict]: + # get prompt template + prompt_template_config = self.get_prompt_template( + app_mode=app_mode, + provider=model_config.provider, + model=model_config.model, + pre_prompt=pre_prompt, + has_context=context is not None, + query_in_prompt=query is not None, + with_memory_prompt=histories is not None, + ) + + variables = {k: inputs[k] for k in prompt_template_config["custom_variable_keys"] if k in inputs} + + for v in prompt_template_config["special_variable_keys"]: + # support #context#, #query# and #histories# + if v == "#context#": + variables["#context#"] = context or "" + elif v == "#query#": + variables["#query#"] = query or "" + elif v == "#histories#": + variables["#histories#"] = histories or "" + + prompt_template = prompt_template_config["prompt_template"] + prompt = prompt_template.format(variables) + + return prompt, prompt_template_config["prompt_rules"] + + def get_prompt_template( + self, + app_mode: AppMode, + provider: str, + model: str, + pre_prompt: str, + has_context: bool, + query_in_prompt: bool, + with_memory_prompt: bool = False, + ) -> dict: + prompt_rules = self._get_prompt_rule(app_mode=app_mode, provider=provider, model=model) + + custom_variable_keys = [] + special_variable_keys = [] + + prompt = "" + for order in prompt_rules["system_prompt_orders"]: + if order == "context_prompt" and has_context: + prompt += prompt_rules["context_prompt"] + special_variable_keys.append("#context#") + elif order == "pre_prompt" and pre_prompt: + prompt += pre_prompt + "\n" + pre_prompt_template = PromptTemplateParser(template=pre_prompt) + custom_variable_keys = pre_prompt_template.variable_keys + elif order == "histories_prompt" and with_memory_prompt: + prompt += prompt_rules["histories_prompt"] + special_variable_keys.append("#histories#") + + if query_in_prompt: + prompt += prompt_rules.get("query_prompt", "{{#query#}}") + special_variable_keys.append("#query#") + + return { + "prompt_template": PromptTemplateParser(template=prompt), + "custom_variable_keys": custom_variable_keys, + "special_variable_keys": special_variable_keys, + "prompt_rules": prompt_rules, + } + + def _get_chat_model_prompt_messages( + self, + app_mode: AppMode, + pre_prompt: str, + inputs: dict, + query: str, + context: Optional[str], + files: Sequence["File"], + memory: Optional[TokenBufferMemory], + model_config: ModelConfigWithCredentialsEntity, + ) -> tuple[list[PromptMessage], Optional[list[str]]]: + prompt_messages: list[PromptMessage] = [] + + # get prompt + prompt, _ = self.get_prompt_str_and_rules( + app_mode=app_mode, + model_config=model_config, + pre_prompt=pre_prompt, + inputs=inputs, + query=None, + context=context, + ) + + if prompt and query: + prompt_messages.append(SystemPromptMessage(content=prompt)) + + if memory: + prompt_messages = self._append_chat_histories( + memory=memory, + memory_config=MemoryConfig( + window=MemoryConfig.WindowConfig( + enabled=False, + ) + ), + prompt_messages=prompt_messages, + model_config=model_config, + ) + + if query: + prompt_messages.append(self.get_last_user_message(query, files)) + else: + prompt_messages.append(self.get_last_user_message(prompt, files)) + + return prompt_messages, None + + def _get_completion_model_prompt_messages( + self, + app_mode: AppMode, + pre_prompt: str, + inputs: dict, + query: str, + context: Optional[str], + files: Sequence["File"], + memory: Optional[TokenBufferMemory], + model_config: ModelConfigWithCredentialsEntity, + ) -> tuple[list[PromptMessage], Optional[list[str]]]: + # get prompt + prompt, prompt_rules = self.get_prompt_str_and_rules( + app_mode=app_mode, + model_config=model_config, + pre_prompt=pre_prompt, + inputs=inputs, + query=query, + context=context, + ) + + if memory: + tmp_human_message = UserPromptMessage(content=prompt) + + rest_tokens = self._calculate_rest_token([tmp_human_message], model_config) + histories = self._get_history_messages_from_memory( + memory=memory, + memory_config=MemoryConfig( + window=MemoryConfig.WindowConfig( + enabled=False, + ) + ), + max_token_limit=rest_tokens, + human_prefix=prompt_rules.get("human_prefix", "Human"), + ai_prefix=prompt_rules.get("assistant_prefix", "Assistant"), + ) + + # get prompt + prompt, prompt_rules = self.get_prompt_str_and_rules( + app_mode=app_mode, + model_config=model_config, + pre_prompt=pre_prompt, + inputs=inputs, + query=query, + context=context, + histories=histories, + ) + + stops = prompt_rules.get("stops") + if stops is not None and len(stops) == 0: + stops = None + + return [self.get_last_user_message(prompt, files)], stops + + def get_last_user_message(self, prompt: str, files: Sequence["File"]) -> UserPromptMessage: + if files: + prompt_message_contents: list[PromptMessageContent] = [] + prompt_message_contents.append(TextPromptMessageContent(data=prompt)) + for file in files: + prompt_message_contents.append(file_manager.to_prompt_message_content(file)) + + prompt_message = UserPromptMessage(content=prompt_message_contents) + else: + prompt_message = UserPromptMessage(content=prompt) + + return prompt_message + + def _get_prompt_rule(self, app_mode: AppMode, provider: str, model: str) -> dict: + """ + Get simple prompt rule. + :param app_mode: app mode + :param provider: model provider + :param model: model name + :return: + """ + prompt_file_name = self._prompt_file_name(app_mode=app_mode, provider=provider, model=model) + + # Check if the prompt file is already loaded + if prompt_file_name in prompt_file_contents: + return cast(dict, prompt_file_contents[prompt_file_name]) + + # Get the absolute path of the subdirectory + prompt_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "prompt_templates") + json_file_path = os.path.join(prompt_path, f"{prompt_file_name}.json") + + # Open the JSON file and read its content + with open(json_file_path, encoding="utf-8") as json_file: + content = json.load(json_file) + + # Store the content of the prompt file + prompt_file_contents[prompt_file_name] = content + + return cast(dict, content) + + def _prompt_file_name(self, app_mode: AppMode, provider: str, model: str) -> str: + # baichuan + is_baichuan = False + if provider == "baichuan": + is_baichuan = True + else: + baichuan_supported_providers = ["huggingface_hub", "openllm", "xinference"] + if provider in baichuan_supported_providers and "baichuan" in model.lower(): + is_baichuan = True + + if is_baichuan: + if app_mode == AppMode.COMPLETION: + return "baichuan_completion" + else: + return "baichuan_chat" + + # common + if app_mode == AppMode.COMPLETION: + return "common_completion" + else: + return "common_chat" diff --git a/api/core/prompt/utils/__init__.py b/api/core/prompt/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/prompt/utils/extract_thread_messages.py b/api/core/prompt/utils/extract_thread_messages.py new file mode 100644 index 0000000000000000000000000000000000000000..f7aef76c87edc864254e2809f75a26dcf2f1c64a --- /dev/null +++ b/api/core/prompt/utils/extract_thread_messages.py @@ -0,0 +1,24 @@ +from typing import Any + +from constants import UUID_NIL + + +def extract_thread_messages(messages: list[Any]): + thread_messages = [] + next_message = None + + for message in messages: + if not message.parent_message_id: + # If the message is regenerated and does not have a parent message, it is the start of a new thread + thread_messages.append(message) + break + + if not next_message: + thread_messages.append(message) + next_message = message.parent_message_id + else: + if next_message in {message.id, UUID_NIL}: + thread_messages.append(message) + next_message = message.parent_message_id + + return thread_messages diff --git a/api/core/prompt/utils/get_thread_messages_length.py b/api/core/prompt/utils/get_thread_messages_length.py new file mode 100644 index 0000000000000000000000000000000000000000..f49466db6d92f8b19ece1a65396e181c02231608 --- /dev/null +++ b/api/core/prompt/utils/get_thread_messages_length.py @@ -0,0 +1,32 @@ +from core.prompt.utils.extract_thread_messages import extract_thread_messages +from extensions.ext_database import db +from models.model import Message + + +def get_thread_messages_length(conversation_id: str) -> int: + """ + Get the number of thread messages based on the parent message id. + """ + # Fetch all messages related to the conversation + query = ( + db.session.query( + Message.id, + Message.parent_message_id, + Message.answer, + ) + .filter( + Message.conversation_id == conversation_id, + ) + .order_by(Message.created_at.desc()) + ) + + messages = query.all() + + # Extract thread messages + thread_messages = extract_thread_messages(messages) + + # Exclude the newly created message with an empty answer + if thread_messages and not thread_messages[0].answer: + thread_messages.pop(0) + + return len(thread_messages) diff --git a/api/core/prompt/utils/prompt_message_util.py b/api/core/prompt/utils/prompt_message_util.py new file mode 100644 index 0000000000000000000000000000000000000000..2f4e65146131be86afdba39f7897e50a206106e6 --- /dev/null +++ b/api/core/prompt/utils/prompt_message_util.py @@ -0,0 +1,114 @@ +from collections.abc import Sequence +from typing import Any, cast + +from core.model_runtime.entities import ( + AssistantPromptMessage, + AudioPromptMessageContent, + ImagePromptMessageContent, + PromptMessage, + PromptMessageContentType, + PromptMessageRole, + TextPromptMessageContent, +) +from core.prompt.simple_prompt_transform import ModelMode + + +class PromptMessageUtil: + @staticmethod + def prompt_messages_to_prompt_for_saving(model_mode: str, prompt_messages: Sequence[PromptMessage]) -> list[dict]: + """ + Prompt messages to prompt for saving. + :param model_mode: model mode + :param prompt_messages: prompt messages + :return: + """ + prompts = [] + if model_mode == ModelMode.CHAT: + tool_calls = [] + for prompt_message in prompt_messages: + if prompt_message.role == PromptMessageRole.USER: + role = "user" + elif prompt_message.role == PromptMessageRole.ASSISTANT: + role = "assistant" + if isinstance(prompt_message, AssistantPromptMessage): + tool_calls = [ + { + "id": tool_call.id, + "type": "function", + "function": { + "name": tool_call.function.name, + "arguments": tool_call.function.arguments, + }, + } + for tool_call in prompt_message.tool_calls + ] + elif prompt_message.role == PromptMessageRole.SYSTEM: + role = "system" + elif prompt_message.role == PromptMessageRole.TOOL: + role = "tool" + else: + continue + + text = "" + files = [] + if isinstance(prompt_message.content, list): + for content in prompt_message.content: + if isinstance(content, TextPromptMessageContent): + text += content.data + elif isinstance(content, ImagePromptMessageContent): + files.append( + { + "type": "image", + "data": content.data[:10] + "...[TRUNCATED]..." + content.data[-10:], + "detail": content.detail.value, + } + ) + elif isinstance(content, AudioPromptMessageContent): + files.append( + { + "type": "audio", + "data": content.data[:10] + "...[TRUNCATED]..." + content.data[-10:], + "format": content.format, + } + ) + else: + text = cast(str, prompt_message.content) + + prompt = {"role": role, "text": text, "files": files} + + if tool_calls: + prompt["tool_calls"] = tool_calls + + prompts.append(prompt) + else: + prompt_message = prompt_messages[0] + text = "" + files = [] + if isinstance(prompt_message.content, list): + for content in prompt_message.content: + if content.type == PromptMessageContentType.TEXT: + content = cast(TextPromptMessageContent, content) + text += content.data + else: + content = cast(ImagePromptMessageContent, content) + files.append( + { + "type": "image", + "data": content.data[:10] + "...[TRUNCATED]..." + content.data[-10:], + "detail": content.detail.value, + } + ) + else: + text = cast(str, prompt_message.content) + + params: dict[str, Any] = { + "role": "user", + "text": text, + } + + if files: + params["files"] = files + + prompts.append(params) + + return prompts diff --git a/api/core/prompt/utils/prompt_template_parser.py b/api/core/prompt/utils/prompt_template_parser.py new file mode 100644 index 0000000000000000000000000000000000000000..8e40674bc193e027607f816e2767c627e8157734 --- /dev/null +++ b/api/core/prompt/utils/prompt_template_parser.py @@ -0,0 +1,46 @@ +import re +from collections.abc import Mapping + +REGEX = re.compile(r"\{\{([a-zA-Z_][a-zA-Z0-9_]{0,29}|#histories#|#query#|#context#)\}\}") +WITH_VARIABLE_TMPL_REGEX = re.compile( + r"\{\{([a-zA-Z_][a-zA-Z0-9_]{0,29}|#[a-zA-Z0-9_]{1,50}\.[a-zA-Z0-9_\.]{1,100}#|#histories#|#query#|#context#)\}\}" +) + + +class PromptTemplateParser: + """ + Rules: + + 1. Template variables must be enclosed in `{{}}`. + 2. The template variable Key can only be: letters + numbers + underscore, with a maximum length of 16 characters, + and can only start with letters and underscores. + 3. The template variable Key cannot contain new lines or spaces, and must comply with rule 2. + 4. In addition to the above, 3 types of special template variable Keys are accepted: + `{{#histories#}}` `{{#query#}}` `{{#context#}}`. No other `{{##}}` template variables are allowed. + """ + + def __init__(self, template: str, with_variable_tmpl: bool = False): + self.template = template + self.with_variable_tmpl = with_variable_tmpl + self.regex = WITH_VARIABLE_TMPL_REGEX if with_variable_tmpl else REGEX + self.variable_keys = self.extract() + + def extract(self) -> list: + # Regular expression to match the template rules + return re.findall(self.regex, self.template) + + def format(self, inputs: Mapping[str, str], remove_template_variables: bool = True) -> str: + def replacer(match): + key = match.group(1) + value = inputs.get(key, match.group(0)) # return original matched string if key not found + + if remove_template_variables and isinstance(value, str): + return PromptTemplateParser.remove_template_variables(value, self.with_variable_tmpl) + return value + + prompt = re.sub(self.regex, replacer, self.template) + return re.sub(r"<\|.*?\|>", "", prompt) + + @classmethod + def remove_template_variables(cls, text: str, with_variable_tmpl: bool = False): + return re.sub(WITH_VARIABLE_TMPL_REGEX if with_variable_tmpl else REGEX, r"{\1}", text) diff --git a/api/core/provider_manager.py b/api/core/provider_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..2430d598ffcb7d4f2871980f9fb6207d4674161b --- /dev/null +++ b/api/core/provider_manager.py @@ -0,0 +1,936 @@ +import json +from collections import defaultdict +from json import JSONDecodeError +from typing import Optional, cast + +from sqlalchemy.exc import IntegrityError + +from configs import dify_config +from core.entities.model_entities import DefaultModelEntity, DefaultModelProviderEntity +from core.entities.provider_configuration import ProviderConfiguration, ProviderConfigurations, ProviderModelBundle +from core.entities.provider_entities import ( + CustomConfiguration, + CustomModelConfiguration, + CustomProviderConfiguration, + ModelLoadBalancingConfiguration, + ModelSettings, + QuotaConfiguration, + QuotaUnit, + SystemConfiguration, +) +from core.helper import encrypter +from core.helper.model_provider_cache import ProviderCredentialsCache, ProviderCredentialsCacheType +from core.helper.position_helper import is_filtered +from core.model_runtime.entities.model_entities import ModelType +from core.model_runtime.entities.provider_entities import ( + ConfigurateMethod, + CredentialFormSchema, + FormType, + ProviderEntity, +) +from core.model_runtime.model_providers import model_provider_factory +from extensions import ext_hosting_provider +from extensions.ext_database import db +from extensions.ext_redis import redis_client +from models.provider import ( + LoadBalancingModelConfig, + Provider, + ProviderModel, + ProviderModelSetting, + ProviderQuotaType, + ProviderType, + TenantDefaultModel, + TenantPreferredModelProvider, +) +from services.feature_service import FeatureService + + +class ProviderManager: + """ + ProviderManager is a class that manages the model providers includes Hosting and Customize Model Providers. + """ + + def __init__(self) -> None: + self.decoding_rsa_key = None + self.decoding_cipher_rsa = None + + def get_configurations(self, tenant_id: str) -> ProviderConfigurations: + """ + Get model provider configurations. + + Construct ProviderConfiguration objects for each provider + Including: + 1. Basic information of the provider + 2. Hosting configuration information, including: + (1. Whether to enable (support) hosting type, if enabled, the following information exists + (2. List of hosting type provider configurations + (including quota type, quota limit, current remaining quota, etc.) + (3. The current hosting type in use (whether there is a quota or not) + paid quotas > provider free quotas > hosting trial quotas + (4. Unified credentials for hosting providers + 3. Custom configuration information, including: + (1. Whether to enable (support) custom type, if enabled, the following information exists + (2. Custom provider configuration (including credentials) + (3. List of custom provider model configurations (including credentials) + 4. Hosting/custom preferred provider type. + Provide methods: + - Get the current configuration (including credentials) + - Get the availability and status of the hosting configuration: active available, + quota_exceeded insufficient quota, unsupported hosting + - Get the availability of custom configuration + Custom provider available conditions: + (1. custom provider credentials available + (2. at least one custom model credentials available + - Verify, update, and delete custom provider configuration + - Verify, update, and delete custom provider model configuration + - Get the list of available models (optional provider filtering, model type filtering) + Append custom provider models to the list + - Get provider instance + - Switch selection priority + + :param tenant_id: + :return: + """ + # Get all provider records of the workspace + provider_name_to_provider_records_dict = self._get_all_providers(tenant_id) + + # Initialize trial provider records if not exist + provider_name_to_provider_records_dict = self._init_trial_provider_records( + tenant_id, provider_name_to_provider_records_dict + ) + + # Get all provider model records of the workspace + provider_name_to_provider_model_records_dict = self._get_all_provider_models(tenant_id) + + # Get all provider entities + provider_entities = model_provider_factory.get_providers() + + # Get All preferred provider types of the workspace + provider_name_to_preferred_model_provider_records_dict = self._get_all_preferred_model_providers(tenant_id) + + # Get All provider model settings + provider_name_to_provider_model_settings_dict = self._get_all_provider_model_settings(tenant_id) + + # Get All load balancing configs + provider_name_to_provider_load_balancing_model_configs_dict = self._get_all_provider_load_balancing_configs( + tenant_id + ) + + provider_configurations = ProviderConfigurations(tenant_id=tenant_id) + + # Construct ProviderConfiguration objects for each provider + for provider_entity in provider_entities: + # handle include, exclude + if is_filtered( + include_set=cast(set[str], dify_config.POSITION_PROVIDER_INCLUDES_SET), + exclude_set=cast(set[str], dify_config.POSITION_PROVIDER_EXCLUDES_SET), + data=provider_entity, + name_func=lambda x: x.provider, + ): + continue + + provider_name = provider_entity.provider + provider_records = provider_name_to_provider_records_dict.get(provider_entity.provider, []) + provider_model_records = provider_name_to_provider_model_records_dict.get(provider_entity.provider, []) + + # Convert to custom configuration + custom_configuration = self._to_custom_configuration( + tenant_id, provider_entity, provider_records, provider_model_records + ) + + # Convert to system configuration + system_configuration = self._to_system_configuration(tenant_id, provider_entity, provider_records) + + # Get preferred provider type + preferred_provider_type_record = provider_name_to_preferred_model_provider_records_dict.get(provider_name) + + if preferred_provider_type_record: + preferred_provider_type = ProviderType.value_of(preferred_provider_type_record.preferred_provider_type) + elif custom_configuration.provider or custom_configuration.models: + preferred_provider_type = ProviderType.CUSTOM + elif system_configuration.enabled: + preferred_provider_type = ProviderType.SYSTEM + else: + preferred_provider_type = ProviderType.CUSTOM + + using_provider_type = preferred_provider_type + has_valid_quota = any(quota_conf.is_valid for quota_conf in system_configuration.quota_configurations) + + if preferred_provider_type == ProviderType.SYSTEM: + if not system_configuration.enabled or not has_valid_quota: + using_provider_type = ProviderType.CUSTOM + + else: + if not custom_configuration.provider and not custom_configuration.models: + if system_configuration.enabled and has_valid_quota: + using_provider_type = ProviderType.SYSTEM + + # Get provider load balancing configs + provider_model_settings = provider_name_to_provider_model_settings_dict.get(provider_name) + + # Get provider load balancing configs + provider_load_balancing_configs = provider_name_to_provider_load_balancing_model_configs_dict.get( + provider_name + ) + + # Convert to model settings + model_settings = self._to_model_settings( + provider_entity=provider_entity, + provider_model_settings=provider_model_settings, + load_balancing_model_configs=provider_load_balancing_configs, + ) + + provider_configuration = ProviderConfiguration( + tenant_id=tenant_id, + provider=provider_entity, + preferred_provider_type=preferred_provider_type, + using_provider_type=using_provider_type, + system_configuration=system_configuration, + custom_configuration=custom_configuration, + model_settings=model_settings, + ) + + provider_configurations[provider_name] = provider_configuration + + # Return the encapsulated object + return provider_configurations + + def get_provider_model_bundle(self, tenant_id: str, provider: str, model_type: ModelType) -> ProviderModelBundle: + """ + Get provider model bundle. + :param tenant_id: workspace id + :param provider: provider name + :param model_type: model type + :return: + """ + provider_configurations = self.get_configurations(tenant_id) + + # get provider instance + provider_configuration = provider_configurations.get(provider) + if not provider_configuration: + raise ValueError(f"Provider {provider} does not exist.") + + provider_instance = provider_configuration.get_provider_instance() + model_type_instance = provider_instance.get_model_instance(model_type) + + return ProviderModelBundle( + configuration=provider_configuration, + provider_instance=provider_instance, + model_type_instance=model_type_instance, + ) + + def get_default_model(self, tenant_id: str, model_type: ModelType) -> Optional[DefaultModelEntity]: + """ + Get default model. + + :param tenant_id: workspace id + :param model_type: model type + :return: + """ + # Get the corresponding TenantDefaultModel record + default_model = ( + db.session.query(TenantDefaultModel) + .filter( + TenantDefaultModel.tenant_id == tenant_id, + TenantDefaultModel.model_type == model_type.to_origin_model_type(), + ) + .first() + ) + + # If it does not exist, get the first available provider model from get_configurations + # and update the TenantDefaultModel record + if not default_model: + # Get provider configurations + provider_configurations = self.get_configurations(tenant_id) + + # get available models from provider_configurations + available_models = provider_configurations.get_models(model_type=model_type, only_active=True) + + if available_models: + available_model = next( + (model for model in available_models if model.model == "gpt-4"), available_models[0] + ) + + default_model = TenantDefaultModel( + tenant_id=tenant_id, + model_type=model_type.to_origin_model_type(), + provider_name=available_model.provider.provider, + model_name=available_model.model, + ) + db.session.add(default_model) + db.session.commit() + + if not default_model: + return None + + provider_instance = model_provider_factory.get_provider_instance(default_model.provider_name) + provider_schema = provider_instance.get_provider_schema() + + return DefaultModelEntity( + model=default_model.model_name, + model_type=model_type, + provider=DefaultModelProviderEntity( + provider=provider_schema.provider, + label=provider_schema.label, + icon_small=provider_schema.icon_small, + icon_large=provider_schema.icon_large, + supported_model_types=provider_schema.supported_model_types, + ), + ) + + def get_first_provider_first_model(self, tenant_id: str, model_type: ModelType) -> tuple[str, str]: + """ + Get names of first model and its provider + + :param tenant_id: workspace id + :param model_type: model type + :return: provider name, model name + """ + provider_configurations = self.get_configurations(tenant_id) + + # get available models from provider_configurations + all_models = provider_configurations.get_models(model_type=model_type, only_active=False) + + return all_models[0].provider.provider, all_models[0].model + + def update_default_model_record( + self, tenant_id: str, model_type: ModelType, provider: str, model: str + ) -> TenantDefaultModel: + """ + Update default model record. + + :param tenant_id: workspace id + :param model_type: model type + :param provider: provider name + :param model: model name + :return: + """ + provider_configurations = self.get_configurations(tenant_id) + if provider not in provider_configurations: + raise ValueError(f"Provider {provider} does not exist.") + + # get available models from provider_configurations + available_models = provider_configurations.get_models(model_type=model_type, only_active=True) + + # check if the model is exist in available models + model_names = [model.model for model in available_models] + if model not in model_names: + raise ValueError(f"Model {model} does not exist.") + + # Get the list of available models from get_configurations and check if it is LLM + default_model = ( + db.session.query(TenantDefaultModel) + .filter( + TenantDefaultModel.tenant_id == tenant_id, + TenantDefaultModel.model_type == model_type.to_origin_model_type(), + ) + .first() + ) + + # create or update TenantDefaultModel record + if default_model: + # update default model + default_model.provider_name = provider + default_model.model_name = model + db.session.commit() + else: + # create default model + default_model = TenantDefaultModel( + tenant_id=tenant_id, + model_type=model_type.value, + provider_name=provider, + model_name=model, + ) + db.session.add(default_model) + db.session.commit() + + return default_model + + @staticmethod + def _get_all_providers(tenant_id: str) -> dict[str, list[Provider]]: + """ + Get all provider records of the workspace. + + :param tenant_id: workspace id + :return: + """ + providers = db.session.query(Provider).filter(Provider.tenant_id == tenant_id, Provider.is_valid == True).all() + + provider_name_to_provider_records_dict = defaultdict(list) + for provider in providers: + provider_name_to_provider_records_dict[provider.provider_name].append(provider) + + return provider_name_to_provider_records_dict + + @staticmethod + def _get_all_provider_models(tenant_id: str) -> dict[str, list[ProviderModel]]: + """ + Get all provider model records of the workspace. + + :param tenant_id: workspace id + :return: + """ + # Get all provider model records of the workspace + provider_models = ( + db.session.query(ProviderModel) + .filter(ProviderModel.tenant_id == tenant_id, ProviderModel.is_valid == True) + .all() + ) + + provider_name_to_provider_model_records_dict = defaultdict(list) + for provider_model in provider_models: + provider_name_to_provider_model_records_dict[provider_model.provider_name].append(provider_model) + + return provider_name_to_provider_model_records_dict + + @staticmethod + def _get_all_preferred_model_providers(tenant_id: str) -> dict[str, TenantPreferredModelProvider]: + """ + Get All preferred provider types of the workspace. + + :param tenant_id: workspace id + :return: + """ + preferred_provider_types = ( + db.session.query(TenantPreferredModelProvider) + .filter(TenantPreferredModelProvider.tenant_id == tenant_id) + .all() + ) + + provider_name_to_preferred_provider_type_records_dict = { + preferred_provider_type.provider_name: preferred_provider_type + for preferred_provider_type in preferred_provider_types + } + + return provider_name_to_preferred_provider_type_records_dict + + @staticmethod + def _get_all_provider_model_settings(tenant_id: str) -> dict[str, list[ProviderModelSetting]]: + """ + Get All provider model settings of the workspace. + + :param tenant_id: workspace id + :return: + """ + provider_model_settings = ( + db.session.query(ProviderModelSetting).filter(ProviderModelSetting.tenant_id == tenant_id).all() + ) + + provider_name_to_provider_model_settings_dict = defaultdict(list) + for provider_model_setting in provider_model_settings: + ( + provider_name_to_provider_model_settings_dict[provider_model_setting.provider_name].append( + provider_model_setting + ) + ) + + return provider_name_to_provider_model_settings_dict + + @staticmethod + def _get_all_provider_load_balancing_configs(tenant_id: str) -> dict[str, list[LoadBalancingModelConfig]]: + """ + Get All provider load balancing configs of the workspace. + + :param tenant_id: workspace id + :return: + """ + cache_key = f"tenant:{tenant_id}:model_load_balancing_enabled" + cache_result = redis_client.get(cache_key) + if cache_result is None: + model_load_balancing_enabled = FeatureService.get_features(tenant_id).model_load_balancing_enabled + redis_client.setex(cache_key, 120, str(model_load_balancing_enabled)) + else: + cache_result = cache_result.decode("utf-8") + model_load_balancing_enabled = cache_result == "True" + + if not model_load_balancing_enabled: + return {} + + provider_load_balancing_configs = ( + db.session.query(LoadBalancingModelConfig).filter(LoadBalancingModelConfig.tenant_id == tenant_id).all() + ) + + provider_name_to_provider_load_balancing_model_configs_dict = defaultdict(list) + for provider_load_balancing_config in provider_load_balancing_configs: + ( + provider_name_to_provider_load_balancing_model_configs_dict[ + provider_load_balancing_config.provider_name + ].append(provider_load_balancing_config) + ) + + return provider_name_to_provider_load_balancing_model_configs_dict + + @staticmethod + def _init_trial_provider_records( + tenant_id: str, provider_name_to_provider_records_dict: dict[str, list] + ) -> dict[str, list]: + """ + Initialize trial provider records if not exists. + + :param tenant_id: workspace id + :param provider_name_to_provider_records_dict: provider name to provider records dict + :return: + """ + # Get hosting configuration + hosting_configuration = ext_hosting_provider.hosting_configuration + + for provider_name, configuration in hosting_configuration.provider_map.items(): + if not configuration.enabled: + continue + + provider_records = provider_name_to_provider_records_dict.get(provider_name) + if not provider_records: + provider_records = [] + + provider_quota_to_provider_record_dict = {} + for provider_record in provider_records: + if provider_record.provider_type != ProviderType.SYSTEM.value: + continue + + provider_quota_to_provider_record_dict[ProviderQuotaType.value_of(provider_record.quota_type)] = ( + provider_record + ) + + for quota in configuration.quotas: + if quota.quota_type == ProviderQuotaType.TRIAL: + # Init trial provider records if not exists + if ProviderQuotaType.TRIAL not in provider_quota_to_provider_record_dict: + try: + # FIXME ignore the type errork, onyl TrialHostingQuota has limit need to change the logic + provider_record = Provider( + tenant_id=tenant_id, + provider_name=provider_name, + provider_type=ProviderType.SYSTEM.value, + quota_type=ProviderQuotaType.TRIAL.value, + quota_limit=quota.quota_limit, # type: ignore + quota_used=0, + is_valid=True, + ) + db.session.add(provider_record) + db.session.commit() + except IntegrityError: + db.session.rollback() + provider_record = ( + db.session.query(Provider) + .filter( + Provider.tenant_id == tenant_id, + Provider.provider_name == provider_name, + Provider.provider_type == ProviderType.SYSTEM.value, + Provider.quota_type == ProviderQuotaType.TRIAL.value, + ) + .first() + ) + + if provider_record and not provider_record.is_valid: + provider_record.is_valid = True + db.session.commit() + + provider_name_to_provider_records_dict[provider_name].append(provider_record) + + return provider_name_to_provider_records_dict + + def _to_custom_configuration( + self, + tenant_id: str, + provider_entity: ProviderEntity, + provider_records: list[Provider], + provider_model_records: list[ProviderModel], + ) -> CustomConfiguration: + """ + Convert to custom configuration. + + :param tenant_id: workspace id + :param provider_entity: provider entity + :param provider_records: provider records + :param provider_model_records: provider model records + :return: + """ + # Get provider credential secret variables + provider_credential_secret_variables = self._extract_secret_variables( + provider_entity.provider_credential_schema.credential_form_schemas + if provider_entity.provider_credential_schema + else [] + ) + + # Get custom provider record + custom_provider_record = None + for provider_record in provider_records: + if provider_record.provider_type == ProviderType.SYSTEM.value: + continue + + if not provider_record.encrypted_config: + continue + + custom_provider_record = provider_record + + # Get custom provider credentials + custom_provider_configuration = None + if custom_provider_record: + provider_credentials_cache = ProviderCredentialsCache( + tenant_id=tenant_id, + identity_id=custom_provider_record.id, + cache_type=ProviderCredentialsCacheType.PROVIDER, + ) + + # Get cached provider credentials + cached_provider_credentials = provider_credentials_cache.get() + + if not cached_provider_credentials: + try: + # fix origin data + if ( + custom_provider_record.encrypted_config + and not custom_provider_record.encrypted_config.startswith("{") + ): + provider_credentials = {"openai_api_key": custom_provider_record.encrypted_config} + else: + provider_credentials = json.loads(custom_provider_record.encrypted_config) + except JSONDecodeError: + provider_credentials = {} + + # Get decoding rsa key and cipher for decrypting credentials + if self.decoding_rsa_key is None or self.decoding_cipher_rsa is None: + self.decoding_rsa_key, self.decoding_cipher_rsa = encrypter.get_decrypt_decoding(tenant_id) + + for variable in provider_credential_secret_variables: + if variable in provider_credentials: + try: + provider_credentials[variable] = encrypter.decrypt_token_with_decoding( + provider_credentials.get(variable) or "", # type: ignore + self.decoding_rsa_key, + self.decoding_cipher_rsa, + ) + except ValueError: + pass + + # cache provider credentials + provider_credentials_cache.set(credentials=provider_credentials) + else: + provider_credentials = cached_provider_credentials + + custom_provider_configuration = CustomProviderConfiguration(credentials=provider_credentials) + + # Get provider model credential secret variables + model_credential_secret_variables = self._extract_secret_variables( + provider_entity.model_credential_schema.credential_form_schemas + if provider_entity.model_credential_schema + else [] + ) + + # Get custom provider model credentials + custom_model_configurations = [] + for provider_model_record in provider_model_records: + if not provider_model_record.encrypted_config: + continue + + provider_model_credentials_cache = ProviderCredentialsCache( + tenant_id=tenant_id, identity_id=provider_model_record.id, cache_type=ProviderCredentialsCacheType.MODEL + ) + + # Get cached provider model credentials + cached_provider_model_credentials = provider_model_credentials_cache.get() + + if not cached_provider_model_credentials: + try: + provider_model_credentials = json.loads(provider_model_record.encrypted_config) + except JSONDecodeError: + continue + + # Get decoding rsa key and cipher for decrypting credentials + if self.decoding_rsa_key is None or self.decoding_cipher_rsa is None: + self.decoding_rsa_key, self.decoding_cipher_rsa = encrypter.get_decrypt_decoding(tenant_id) + + for variable in model_credential_secret_variables: + if variable in provider_model_credentials: + try: + provider_model_credentials[variable] = encrypter.decrypt_token_with_decoding( + provider_model_credentials.get(variable), + self.decoding_rsa_key, + self.decoding_cipher_rsa, + ) + except ValueError: + pass + + # cache provider model credentials + provider_model_credentials_cache.set(credentials=provider_model_credentials) + else: + provider_model_credentials = cached_provider_model_credentials + + custom_model_configurations.append( + CustomModelConfiguration( + model=provider_model_record.model_name, + model_type=ModelType.value_of(provider_model_record.model_type), + credentials=provider_model_credentials, + ) + ) + + return CustomConfiguration(provider=custom_provider_configuration, models=custom_model_configurations) + + def _to_system_configuration( + self, tenant_id: str, provider_entity: ProviderEntity, provider_records: list[Provider] + ) -> SystemConfiguration: + """ + Convert to system configuration. + + :param tenant_id: workspace id + :param provider_entity: provider entity + :param provider_records: provider records + :return: + """ + # Get hosting configuration + hosting_configuration = ext_hosting_provider.hosting_configuration + + provider_hosting_configuration = hosting_configuration.provider_map.get(provider_entity.provider) + if provider_hosting_configuration is None or not provider_hosting_configuration.enabled: + return SystemConfiguration(enabled=False) + + # Convert provider_records to dict + quota_type_to_provider_records_dict = {} + for provider_record in provider_records: + if provider_record.provider_type != ProviderType.SYSTEM.value: + continue + + quota_type_to_provider_records_dict[ProviderQuotaType.value_of(provider_record.quota_type)] = ( + provider_record + ) + quota_configurations = [] + for provider_quota in provider_hosting_configuration.quotas: + if provider_quota.quota_type not in quota_type_to_provider_records_dict: + if provider_quota.quota_type == ProviderQuotaType.FREE: + quota_configuration = QuotaConfiguration( + quota_type=provider_quota.quota_type, + quota_unit=provider_hosting_configuration.quota_unit or QuotaUnit.TOKENS, + quota_used=0, + quota_limit=0, + is_valid=False, + restrict_models=provider_quota.restrict_models, + ) + else: + continue + else: + provider_record = quota_type_to_provider_records_dict[provider_quota.quota_type] + + quota_configuration = QuotaConfiguration( + quota_type=provider_quota.quota_type, + quota_unit=provider_hosting_configuration.quota_unit or QuotaUnit.TOKENS, + quota_used=provider_record.quota_used, + quota_limit=provider_record.quota_limit, + is_valid=provider_record.quota_limit > provider_record.quota_used + or provider_record.quota_limit == -1, + restrict_models=provider_quota.restrict_models, + ) + + quota_configurations.append(quota_configuration) + + if len(quota_configurations) == 0: + return SystemConfiguration(enabled=False) + + current_quota_type = self._choice_current_using_quota_type(quota_configurations) + + current_using_credentials = provider_hosting_configuration.credentials + if current_quota_type == ProviderQuotaType.FREE: + provider_record_quota_free = quota_type_to_provider_records_dict.get(current_quota_type) + + if provider_record_quota_free: + provider_credentials_cache = ProviderCredentialsCache( + tenant_id=tenant_id, + identity_id=provider_record_quota_free.id, + cache_type=ProviderCredentialsCacheType.PROVIDER, + ) + + # Get cached provider credentials + cached_provider_credentials = provider_credentials_cache.get() + + if not cached_provider_credentials: + try: + provider_credentials = json.loads(provider_record.encrypted_config) + except JSONDecodeError: + provider_credentials = {} + + # Get provider credential secret variables + provider_credential_secret_variables = self._extract_secret_variables( + provider_entity.provider_credential_schema.credential_form_schemas + if provider_entity.provider_credential_schema + else [] + ) + + # Get decoding rsa key and cipher for decrypting credentials + if self.decoding_rsa_key is None or self.decoding_cipher_rsa is None: + self.decoding_rsa_key, self.decoding_cipher_rsa = encrypter.get_decrypt_decoding(tenant_id) + + for variable in provider_credential_secret_variables: + if variable in provider_credentials: + try: + provider_credentials[variable] = encrypter.decrypt_token_with_decoding( + provider_credentials.get(variable), self.decoding_rsa_key, self.decoding_cipher_rsa + ) + except ValueError: + pass + + current_using_credentials = provider_credentials or {} + + # cache provider credentials + provider_credentials_cache.set(credentials=current_using_credentials) + else: + current_using_credentials = cached_provider_credentials + else: + current_using_credentials = {} + quota_configurations = [] + + return SystemConfiguration( + enabled=True, + current_quota_type=current_quota_type, + quota_configurations=quota_configurations, + credentials=current_using_credentials, + ) + + @staticmethod + def _choice_current_using_quota_type(quota_configurations: list[QuotaConfiguration]) -> ProviderQuotaType: + """ + Choice current using quota type. + paid quotas > provider free quotas > hosting trial quotas + If there is still quota for the corresponding quota type according to the sorting, + + :param quota_configurations: + :return: + """ + # convert to dict + quota_type_to_quota_configuration_dict = { + quota_configuration.quota_type: quota_configuration for quota_configuration in quota_configurations + } + + last_quota_configuration = None + for quota_type in [ProviderQuotaType.PAID, ProviderQuotaType.FREE, ProviderQuotaType.TRIAL]: + if quota_type in quota_type_to_quota_configuration_dict: + last_quota_configuration = quota_type_to_quota_configuration_dict[quota_type] + if last_quota_configuration.is_valid: + return quota_type + + if last_quota_configuration: + return last_quota_configuration.quota_type + + raise ValueError("No quota type available") + + @staticmethod + def _extract_secret_variables(credential_form_schemas: list[CredentialFormSchema]) -> list[str]: + """ + Extract secret input form variables. + + :param credential_form_schemas: + :return: + """ + secret_input_form_variables = [] + for credential_form_schema in credential_form_schemas: + if credential_form_schema.type == FormType.SECRET_INPUT: + secret_input_form_variables.append(credential_form_schema.variable) + + return secret_input_form_variables + + def _to_model_settings( + self, + provider_entity: ProviderEntity, + provider_model_settings: Optional[list[ProviderModelSetting]] = None, + load_balancing_model_configs: Optional[list[LoadBalancingModelConfig]] = None, + ) -> list[ModelSettings]: + """ + Convert to model settings. + :param provider_entity: provider entity + :param provider_model_settings: provider model settings include enabled, load balancing enabled + :param load_balancing_model_configs: load balancing model configs + :return: + """ + # Get provider model credential secret variables + if ConfigurateMethod.PREDEFINED_MODEL in provider_entity.configurate_methods: + model_credential_secret_variables = self._extract_secret_variables( + provider_entity.provider_credential_schema.credential_form_schemas + if provider_entity.provider_credential_schema + else [] + ) + else: + model_credential_secret_variables = self._extract_secret_variables( + provider_entity.model_credential_schema.credential_form_schemas + if provider_entity.model_credential_schema + else [] + ) + + model_settings: list[ModelSettings] = [] + if not provider_model_settings: + return model_settings + + for provider_model_setting in provider_model_settings: + load_balancing_configs = [] + if provider_model_setting.load_balancing_enabled and load_balancing_model_configs: + for load_balancing_model_config in load_balancing_model_configs: + if ( + load_balancing_model_config.model_name == provider_model_setting.model_name + and load_balancing_model_config.model_type == provider_model_setting.model_type + ): + if not load_balancing_model_config.enabled: + continue + + if not load_balancing_model_config.encrypted_config: + if load_balancing_model_config.name == "__inherit__": + load_balancing_configs.append( + ModelLoadBalancingConfiguration( + id=load_balancing_model_config.id, + name=load_balancing_model_config.name, + credentials={}, + ) + ) + continue + + provider_model_credentials_cache = ProviderCredentialsCache( + tenant_id=load_balancing_model_config.tenant_id, + identity_id=load_balancing_model_config.id, + cache_type=ProviderCredentialsCacheType.LOAD_BALANCING_MODEL, + ) + + # Get cached provider model credentials + cached_provider_model_credentials = provider_model_credentials_cache.get() + + if not cached_provider_model_credentials: + try: + provider_model_credentials = json.loads(load_balancing_model_config.encrypted_config) + except JSONDecodeError: + continue + + # Get decoding rsa key and cipher for decrypting credentials + if self.decoding_rsa_key is None or self.decoding_cipher_rsa is None: + self.decoding_rsa_key, self.decoding_cipher_rsa = encrypter.get_decrypt_decoding( + load_balancing_model_config.tenant_id + ) + + for variable in model_credential_secret_variables: + if variable in provider_model_credentials: + try: + provider_model_credentials[variable] = encrypter.decrypt_token_with_decoding( + provider_model_credentials.get(variable), + self.decoding_rsa_key, + self.decoding_cipher_rsa, + ) + except ValueError: + pass + + # cache provider model credentials + provider_model_credentials_cache.set(credentials=provider_model_credentials) + else: + provider_model_credentials = cached_provider_model_credentials + + load_balancing_configs.append( + ModelLoadBalancingConfiguration( + id=load_balancing_model_config.id, + name=load_balancing_model_config.name, + credentials=provider_model_credentials, + ) + ) + + model_settings.append( + ModelSettings( + model=provider_model_setting.model_name, + model_type=ModelType.value_of(provider_model_setting.model_type), + enabled=provider_model_setting.enabled, + load_balancing_configs=load_balancing_configs if len(load_balancing_configs) > 1 else [], + ) + ) + + return model_settings diff --git a/api/core/rag/__init__.py b/api/core/rag/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/cleaner/clean_processor.py b/api/core/rag/cleaner/clean_processor.py new file mode 100644 index 0000000000000000000000000000000000000000..754b0d18b71af9097fab6f098e55cf58a922c7b2 --- /dev/null +++ b/api/core/rag/cleaner/clean_processor.py @@ -0,0 +1,36 @@ +import re + + +class CleanProcessor: + @classmethod + def clean(cls, text: str, process_rule: dict) -> str: + # default clean + # remove invalid symbol + text = re.sub(r"<\|", "<", text) + text = re.sub(r"\|>", ">", text) + text = re.sub(r"[\x00-\x08\x0B\x0C\x0E-\x1F\x7F\xEF\xBF\xBE]", "", text) + # Unicode U+FFFE + text = re.sub("\ufffe", "", text) + + rules = process_rule["rules"] if process_rule else {} + if "pre_processing_rules" in rules: + pre_processing_rules = rules["pre_processing_rules"] + for pre_processing_rule in pre_processing_rules: + if pre_processing_rule["id"] == "remove_extra_spaces" and pre_processing_rule["enabled"] is True: + # Remove extra spaces + pattern = r"\n{3,}" + text = re.sub(pattern, "\n\n", text) + pattern = r"[\t\f\r\x20\u00a0\u1680\u180e\u2000-\u200a\u202f\u205f\u3000]{2,}" + text = re.sub(pattern, " ", text) + elif pre_processing_rule["id"] == "remove_urls_emails" and pre_processing_rule["enabled"] is True: + # Remove email + pattern = r"([a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+)" + text = re.sub(pattern, "", text) + + # Remove URL + pattern = r"https?://[^\s]+" + text = re.sub(pattern, "", text) + return text + + def filter_string(self, text): + return text diff --git a/api/core/rag/cleaner/cleaner_base.py b/api/core/rag/cleaner/cleaner_base.py new file mode 100644 index 0000000000000000000000000000000000000000..d3bc2f765e96543f6f11a899e0bd6d7f8cb009ee --- /dev/null +++ b/api/core/rag/cleaner/cleaner_base.py @@ -0,0 +1,11 @@ +"""Abstract interface for document cleaner implementations.""" + +from abc import ABC, abstractmethod + + +class BaseCleaner(ABC): + """Interface for clean chunk content.""" + + @abstractmethod + def clean(self, content: str): + raise NotImplementedError diff --git a/api/core/rag/cleaner/unstructured/unstructured_extra_whitespace_cleaner.py b/api/core/rag/cleaner/unstructured/unstructured_extra_whitespace_cleaner.py new file mode 100644 index 0000000000000000000000000000000000000000..167a919e69aa313b36042fb244ba7369ba39c5ef --- /dev/null +++ b/api/core/rag/cleaner/unstructured/unstructured_extra_whitespace_cleaner.py @@ -0,0 +1,12 @@ +"""Abstract interface for document clean implementations.""" + +from core.rag.cleaner.cleaner_base import BaseCleaner + + +class UnstructuredNonAsciiCharsCleaner(BaseCleaner): + def clean(self, content) -> str: + """clean document content.""" + from unstructured.cleaners.core import clean_extra_whitespace + + # Returns "ITEM 1A: RISK FACTORS" + return clean_extra_whitespace(content) diff --git a/api/core/rag/cleaner/unstructured/unstructured_group_broken_paragraphs_cleaner.py b/api/core/rag/cleaner/unstructured/unstructured_group_broken_paragraphs_cleaner.py new file mode 100644 index 0000000000000000000000000000000000000000..9c682d29db376dae66ed33765aab92f490a54eff --- /dev/null +++ b/api/core/rag/cleaner/unstructured/unstructured_group_broken_paragraphs_cleaner.py @@ -0,0 +1,15 @@ +"""Abstract interface for document clean implementations.""" + +from core.rag.cleaner.cleaner_base import BaseCleaner + + +class UnstructuredGroupBrokenParagraphsCleaner(BaseCleaner): + def clean(self, content) -> str: + """clean document content.""" + import re + + from unstructured.cleaners.core import group_broken_paragraphs + + para_split_re = re.compile(r"(\s*\n\s*){3}") + + return group_broken_paragraphs(content, paragraph_split=para_split_re) diff --git a/api/core/rag/cleaner/unstructured/unstructured_non_ascii_chars_cleaner.py b/api/core/rag/cleaner/unstructured/unstructured_non_ascii_chars_cleaner.py new file mode 100644 index 0000000000000000000000000000000000000000..0cdbb171e1081e32406c01e9ab69e7363e27f48b --- /dev/null +++ b/api/core/rag/cleaner/unstructured/unstructured_non_ascii_chars_cleaner.py @@ -0,0 +1,12 @@ +"""Abstract interface for document clean implementations.""" + +from core.rag.cleaner.cleaner_base import BaseCleaner + + +class UnstructuredNonAsciiCharsCleaner(BaseCleaner): + def clean(self, content) -> str: + """clean document content.""" + from unstructured.cleaners.core import clean_non_ascii_chars + + # Returns "This text contains non-ascii characters!" + return clean_non_ascii_chars(content) diff --git a/api/core/rag/cleaner/unstructured/unstructured_replace_unicode_quotes_cleaner.py b/api/core/rag/cleaner/unstructured/unstructured_replace_unicode_quotes_cleaner.py new file mode 100644 index 0000000000000000000000000000000000000000..9f42044a2d5db8daf12137669e09f86fb8795f4b --- /dev/null +++ b/api/core/rag/cleaner/unstructured/unstructured_replace_unicode_quotes_cleaner.py @@ -0,0 +1,12 @@ +"""Abstract interface for document clean implementations.""" + +from core.rag.cleaner.cleaner_base import BaseCleaner + + +class UnstructuredNonAsciiCharsCleaner(BaseCleaner): + def clean(self, content) -> str: + """Replaces unicode quote characters, such as the \x91 character in a string.""" + + from unstructured.cleaners.core import replace_unicode_quotes + + return replace_unicode_quotes(content) diff --git a/api/core/rag/cleaner/unstructured/unstructured_translate_text_cleaner.py b/api/core/rag/cleaner/unstructured/unstructured_translate_text_cleaner.py new file mode 100644 index 0000000000000000000000000000000000000000..32ae7217e878a58e36763a173114bb77939e5cfc --- /dev/null +++ b/api/core/rag/cleaner/unstructured/unstructured_translate_text_cleaner.py @@ -0,0 +1,11 @@ +"""Abstract interface for document clean implementations.""" + +from core.rag.cleaner.cleaner_base import BaseCleaner + + +class UnstructuredTranslateTextCleaner(BaseCleaner): + def clean(self, content) -> str: + """clean document content.""" + from unstructured.cleaners.translate import translate_text + + return translate_text(content) diff --git a/api/core/rag/data_post_processor/__init__.py b/api/core/rag/data_post_processor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/data_post_processor/data_post_processor.py b/api/core/rag/data_post_processor/data_post_processor.py new file mode 100644 index 0000000000000000000000000000000000000000..d17d76333ee70585cbf59f2d93713b937402ec9a --- /dev/null +++ b/api/core/rag/data_post_processor/data_post_processor.py @@ -0,0 +1,99 @@ +from typing import Optional + +from core.model_manager import ModelInstance, ModelManager +from core.model_runtime.entities.model_entities import ModelType +from core.model_runtime.errors.invoke import InvokeAuthorizationError +from core.rag.data_post_processor.reorder import ReorderRunner +from core.rag.models.document import Document +from core.rag.rerank.entity.weight import KeywordSetting, VectorSetting, Weights +from core.rag.rerank.rerank_base import BaseRerankRunner +from core.rag.rerank.rerank_factory import RerankRunnerFactory +from core.rag.rerank.rerank_type import RerankMode + + +class DataPostProcessor: + """Interface for data post-processing document.""" + + def __init__( + self, + tenant_id: str, + reranking_mode: str, + reranking_model: Optional[dict] = None, + weights: Optional[dict] = None, + reorder_enabled: bool = False, + ): + self.rerank_runner = self._get_rerank_runner(reranking_mode, tenant_id, reranking_model, weights) + self.reorder_runner = self._get_reorder_runner(reorder_enabled) + + def invoke( + self, + query: str, + documents: list[Document], + score_threshold: Optional[float] = None, + top_n: Optional[int] = None, + user: Optional[str] = None, + ) -> list[Document]: + if self.rerank_runner: + documents = self.rerank_runner.run(query, documents, score_threshold, top_n, user) + + if self.reorder_runner: + documents = self.reorder_runner.run(documents) + + return documents + + def _get_rerank_runner( + self, + reranking_mode: str, + tenant_id: str, + reranking_model: Optional[dict] = None, + weights: Optional[dict] = None, + ) -> Optional[BaseRerankRunner]: + if reranking_mode == RerankMode.WEIGHTED_SCORE.value and weights: + runner = RerankRunnerFactory.create_rerank_runner( + runner_type=reranking_mode, + tenant_id=tenant_id, + weights=Weights( + vector_setting=VectorSetting( + vector_weight=weights["vector_setting"]["vector_weight"], + embedding_provider_name=weights["vector_setting"]["embedding_provider_name"], + embedding_model_name=weights["vector_setting"]["embedding_model_name"], + ), + keyword_setting=KeywordSetting( + keyword_weight=weights["keyword_setting"]["keyword_weight"], + ), + ), + ) + return runner + elif reranking_mode == RerankMode.RERANKING_MODEL.value: + rerank_model_instance = self._get_rerank_model_instance(tenant_id, reranking_model) + if rerank_model_instance is None: + return None + runner = RerankRunnerFactory.create_rerank_runner( + runner_type=reranking_mode, rerank_model_instance=rerank_model_instance + ) + return runner + return None + + def _get_reorder_runner(self, reorder_enabled) -> Optional[ReorderRunner]: + if reorder_enabled: + return ReorderRunner() + return None + + def _get_rerank_model_instance(self, tenant_id: str, reranking_model: Optional[dict]) -> ModelInstance | None: + if reranking_model: + try: + model_manager = ModelManager() + reranking_provider_name = reranking_model.get("reranking_provider_name") + reranking_model_name = reranking_model.get("reranking_model_name") + if not reranking_provider_name or not reranking_model_name: + return None + rerank_model_instance = model_manager.get_model_instance( + tenant_id=tenant_id, + provider=reranking_provider_name, + model_type=ModelType.RERANK, + model=reranking_model_name, + ) + return rerank_model_instance + except InvokeAuthorizationError: + return None + return None diff --git a/api/core/rag/data_post_processor/reorder.py b/api/core/rag/data_post_processor/reorder.py new file mode 100644 index 0000000000000000000000000000000000000000..a9a0885241e4fd13b584e6982992d8082b7fb797 --- /dev/null +++ b/api/core/rag/data_post_processor/reorder.py @@ -0,0 +1,17 @@ +from core.rag.models.document import Document + + +class ReorderRunner: + def run(self, documents: list[Document]) -> list[Document]: + # Retrieve elements from odd indices (0, 2, 4, etc.) of the documents list + odd_elements = documents[::2] + + # Retrieve elements from even indices (1, 3, 5, etc.) of the documents list + even_elements = documents[1::2] + + # Reverse the list of elements from even indices + even_elements_reversed = even_elements[::-1] + + new_documents = odd_elements + even_elements_reversed + + return new_documents diff --git a/api/core/rag/datasource/__init__.py b/api/core/rag/datasource/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/keyword/__init__.py b/api/core/rag/datasource/keyword/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/keyword/jieba/__init__.py b/api/core/rag/datasource/keyword/jieba/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/keyword/jieba/jieba.py b/api/core/rag/datasource/keyword/jieba/jieba.py new file mode 100644 index 0000000000000000000000000000000000000000..95a2316f1da4dd58b5b8107b95f8346742c5cc0e --- /dev/null +++ b/api/core/rag/datasource/keyword/jieba/jieba.py @@ -0,0 +1,258 @@ +import json +from collections import defaultdict +from typing import Any, Optional + +from pydantic import BaseModel + +from configs import dify_config +from core.rag.datasource.keyword.jieba.jieba_keyword_table_handler import JiebaKeywordTableHandler +from core.rag.datasource.keyword.keyword_base import BaseKeyword +from core.rag.models.document import Document +from extensions.ext_database import db +from extensions.ext_redis import redis_client +from extensions.ext_storage import storage +from models.dataset import Dataset, DatasetKeywordTable, DocumentSegment + + +class KeywordTableConfig(BaseModel): + max_keywords_per_chunk: int = 10 + + +class Jieba(BaseKeyword): + def __init__(self, dataset: Dataset): + super().__init__(dataset) + self._config = KeywordTableConfig() + + def create(self, texts: list[Document], **kwargs) -> BaseKeyword: + lock_name = "keyword_indexing_lock_{}".format(self.dataset.id) + with redis_client.lock(lock_name, timeout=600): + keyword_table_handler = JiebaKeywordTableHandler() + keyword_table = self._get_dataset_keyword_table() + for text in texts: + keywords = keyword_table_handler.extract_keywords( + text.page_content, self._config.max_keywords_per_chunk + ) + if text.metadata is not None: + self._update_segment_keywords(self.dataset.id, text.metadata["doc_id"], list(keywords)) + keyword_table = self._add_text_to_keyword_table( + keyword_table or {}, text.metadata["doc_id"], list(keywords) + ) + + self._save_dataset_keyword_table(keyword_table) + + return self + + def add_texts(self, texts: list[Document], **kwargs): + lock_name = "keyword_indexing_lock_{}".format(self.dataset.id) + with redis_client.lock(lock_name, timeout=600): + keyword_table_handler = JiebaKeywordTableHandler() + + keyword_table = self._get_dataset_keyword_table() + keywords_list = kwargs.get("keywords_list") + for i in range(len(texts)): + text = texts[i] + if keywords_list: + keywords = keywords_list[i] + if not keywords: + keywords = keyword_table_handler.extract_keywords( + text.page_content, self._config.max_keywords_per_chunk + ) + else: + keywords = keyword_table_handler.extract_keywords( + text.page_content, self._config.max_keywords_per_chunk + ) + if text.metadata is not None: + self._update_segment_keywords(self.dataset.id, text.metadata["doc_id"], list(keywords)) + keyword_table = self._add_text_to_keyword_table( + keyword_table or {}, text.metadata["doc_id"], list(keywords) + ) + + self._save_dataset_keyword_table(keyword_table) + + def text_exists(self, id: str) -> bool: + keyword_table = self._get_dataset_keyword_table() + if keyword_table is None: + return False + return id in set.union(*keyword_table.values()) + + def delete_by_ids(self, ids: list[str]) -> None: + lock_name = "keyword_indexing_lock_{}".format(self.dataset.id) + with redis_client.lock(lock_name, timeout=600): + keyword_table = self._get_dataset_keyword_table() + if keyword_table is not None: + keyword_table = self._delete_ids_from_keyword_table(keyword_table, ids) + + self._save_dataset_keyword_table(keyword_table) + + def search(self, query: str, **kwargs: Any) -> list[Document]: + keyword_table = self._get_dataset_keyword_table() + + k = kwargs.get("top_k", 4) + + sorted_chunk_indices = self._retrieve_ids_by_query(keyword_table or {}, query, k) + + documents = [] + for chunk_index in sorted_chunk_indices: + segment = ( + db.session.query(DocumentSegment) + .filter(DocumentSegment.dataset_id == self.dataset.id, DocumentSegment.index_node_id == chunk_index) + .first() + ) + + if segment: + documents.append( + Document( + page_content=segment.content, + metadata={ + "doc_id": chunk_index, + "doc_hash": segment.index_node_hash, + "document_id": segment.document_id, + "dataset_id": segment.dataset_id, + }, + ) + ) + + return documents + + def delete(self) -> None: + lock_name = "keyword_indexing_lock_{}".format(self.dataset.id) + with redis_client.lock(lock_name, timeout=600): + dataset_keyword_table = self.dataset.dataset_keyword_table + if dataset_keyword_table: + db.session.delete(dataset_keyword_table) + db.session.commit() + if dataset_keyword_table.data_source_type != "database": + file_key = "keyword_files/" + self.dataset.tenant_id + "/" + self.dataset.id + ".txt" + storage.delete(file_key) + + def _save_dataset_keyword_table(self, keyword_table): + keyword_table_dict = { + "__type__": "keyword_table", + "__data__": {"index_id": self.dataset.id, "summary": None, "table": keyword_table}, + } + dataset_keyword_table = self.dataset.dataset_keyword_table + keyword_data_source_type = dataset_keyword_table.data_source_type + if keyword_data_source_type == "database": + dataset_keyword_table.keyword_table = json.dumps(keyword_table_dict, cls=SetEncoder) + db.session.commit() + else: + file_key = "keyword_files/" + self.dataset.tenant_id + "/" + self.dataset.id + ".txt" + if storage.exists(file_key): + storage.delete(file_key) + storage.save(file_key, json.dumps(keyword_table_dict, cls=SetEncoder).encode("utf-8")) + + def _get_dataset_keyword_table(self) -> Optional[dict]: + dataset_keyword_table = self.dataset.dataset_keyword_table + if dataset_keyword_table: + keyword_table_dict = dataset_keyword_table.keyword_table_dict + if keyword_table_dict: + return dict(keyword_table_dict["__data__"]["table"]) + else: + keyword_data_source_type = dify_config.KEYWORD_DATA_SOURCE_TYPE + dataset_keyword_table = DatasetKeywordTable( + dataset_id=self.dataset.id, + keyword_table="", + data_source_type=keyword_data_source_type, + ) + if keyword_data_source_type == "database": + dataset_keyword_table.keyword_table = json.dumps( + { + "__type__": "keyword_table", + "__data__": {"index_id": self.dataset.id, "summary": None, "table": {}}, + }, + cls=SetEncoder, + ) + db.session.add(dataset_keyword_table) + db.session.commit() + + return {} + + def _add_text_to_keyword_table(self, keyword_table: dict, id: str, keywords: list[str]) -> dict: + for keyword in keywords: + if keyword not in keyword_table: + keyword_table[keyword] = set() + keyword_table[keyword].add(id) + return keyword_table + + def _delete_ids_from_keyword_table(self, keyword_table: dict, ids: list[str]) -> dict: + # get set of ids that correspond to node + node_idxs_to_delete = set(ids) + + # delete node_idxs from keyword to node idxs mapping + keywords_to_delete = set() + for keyword, node_idxs in keyword_table.items(): + if node_idxs_to_delete.intersection(node_idxs): + keyword_table[keyword] = node_idxs.difference(node_idxs_to_delete) + if not keyword_table[keyword]: + keywords_to_delete.add(keyword) + + for keyword in keywords_to_delete: + del keyword_table[keyword] + + return keyword_table + + def _retrieve_ids_by_query(self, keyword_table: dict, query: str, k: int = 4): + keyword_table_handler = JiebaKeywordTableHandler() + keywords = keyword_table_handler.extract_keywords(query) + + # go through text chunks in order of most matching keywords + chunk_indices_count: dict[str, int] = defaultdict(int) + keywords_list = [keyword for keyword in keywords if keyword in set(keyword_table.keys())] + for keyword in keywords_list: + for node_id in keyword_table[keyword]: + chunk_indices_count[node_id] += 1 + + sorted_chunk_indices = sorted( + chunk_indices_count.keys(), + key=lambda x: chunk_indices_count[x], + reverse=True, + ) + + return sorted_chunk_indices[:k] + + def _update_segment_keywords(self, dataset_id: str, node_id: str, keywords: list[str]): + document_segment = ( + db.session.query(DocumentSegment) + .filter(DocumentSegment.dataset_id == dataset_id, DocumentSegment.index_node_id == node_id) + .first() + ) + if document_segment: + document_segment.keywords = keywords + db.session.add(document_segment) + db.session.commit() + + def create_segment_keywords(self, node_id: str, keywords: list[str]): + keyword_table = self._get_dataset_keyword_table() + self._update_segment_keywords(self.dataset.id, node_id, keywords) + keyword_table = self._add_text_to_keyword_table(keyword_table or {}, node_id, keywords) + self._save_dataset_keyword_table(keyword_table) + + def multi_create_segment_keywords(self, pre_segment_data_list: list): + keyword_table_handler = JiebaKeywordTableHandler() + keyword_table = self._get_dataset_keyword_table() + for pre_segment_data in pre_segment_data_list: + segment = pre_segment_data["segment"] + if pre_segment_data["keywords"]: + segment.keywords = pre_segment_data["keywords"] + keyword_table = self._add_text_to_keyword_table( + keyword_table or {}, segment.index_node_id, pre_segment_data["keywords"] + ) + else: + keywords = keyword_table_handler.extract_keywords(segment.content, self._config.max_keywords_per_chunk) + segment.keywords = list(keywords) + keyword_table = self._add_text_to_keyword_table( + keyword_table or {}, segment.index_node_id, list(keywords) + ) + self._save_dataset_keyword_table(keyword_table) + + def update_segment_keywords_index(self, node_id: str, keywords: list[str]): + keyword_table = self._get_dataset_keyword_table() + keyword_table = self._add_text_to_keyword_table(keyword_table or {}, node_id, keywords) + self._save_dataset_keyword_table(keyword_table) + + +class SetEncoder(json.JSONEncoder): + def default(self, obj): + if isinstance(obj, set): + return list(obj) + return super().default(obj) diff --git a/api/core/rag/datasource/keyword/jieba/jieba_keyword_table_handler.py b/api/core/rag/datasource/keyword/jieba/jieba_keyword_table_handler.py new file mode 100644 index 0000000000000000000000000000000000000000..a6214d955b1dddb87035bb3319c75e5dfe98db16 --- /dev/null +++ b/api/core/rag/datasource/keyword/jieba/jieba_keyword_table_handler.py @@ -0,0 +1,37 @@ +import re +from typing import Optional, cast + + +class JiebaKeywordTableHandler: + def __init__(self): + import jieba.analyse # type: ignore + + from core.rag.datasource.keyword.jieba.stopwords import STOPWORDS + + jieba.analyse.default_tfidf.stop_words = STOPWORDS # type: ignore + + def extract_keywords(self, text: str, max_keywords_per_chunk: Optional[int] = 10) -> set[str]: + """Extract keywords with JIEBA tfidf.""" + import jieba.analyse # type: ignore + + keywords = jieba.analyse.extract_tags( + sentence=text, + topK=max_keywords_per_chunk, + ) + # jieba.analyse.extract_tags returns list[Any] when withFlag is False by default. + keywords = cast(list[str], keywords) + + return set(self._expand_tokens_with_subtokens(set(keywords))) + + def _expand_tokens_with_subtokens(self, tokens: set[str]) -> set[str]: + """Get subtokens from a list of tokens., filtering for stopwords.""" + from core.rag.datasource.keyword.jieba.stopwords import STOPWORDS + + results = set() + for token in tokens: + results.add(token) + sub_tokens = re.findall(r"\w+", token) + if len(sub_tokens) > 1: + results.update({w for w in sub_tokens if w not in list(STOPWORDS)}) + + return results diff --git a/api/core/rag/datasource/keyword/jieba/stopwords.py b/api/core/rag/datasource/keyword/jieba/stopwords.py new file mode 100644 index 0000000000000000000000000000000000000000..9abe78d6ef7e8d99009da0d5b4edbfdaef36eb5c --- /dev/null +++ b/api/core/rag/datasource/keyword/jieba/stopwords.py @@ -0,0 +1,1380 @@ +STOPWORDS = { + "during", + "when", + "but", + "then", + "further", + "isn", + "mustn't", + "until", + "own", + "i", + "couldn", + "y", + "only", + "you've", + "ours", + "who", + "where", + "ourselves", + "has", + "to", + "was", + "didn't", + "themselves", + "if", + "against", + "through", + "her", + "an", + "your", + "can", + "those", + "didn", + "about", + "aren't", + "shan't", + "be", + "not", + "these", + "again", + "so", + "t", + "theirs", + "weren", + "won't", + "won", + "itself", + "just", + "same", + "while", + "why", + "doesn", + "aren", + "him", + "haven", + "for", + "you'll", + "that", + "we", + "am", + "d", + "by", + "having", + "wasn't", + "than", + "weren't", + "out", + "from", + "now", + "their", + "too", + "hadn", + "o", + "needn", + "most", + "it", + "under", + "needn't", + "any", + "some", + "few", + "ll", + "hers", + "which", + "m", + "you're", + "off", + "other", + "had", + "she", + "you'd", + "do", + "you", + "does", + "s", + "will", + "each", + "wouldn't", + "hasn't", + "such", + "more", + "whom", + "she's", + "my", + "yours", + "yourself", + "of", + "on", + "very", + "hadn't", + "with", + "yourselves", + "been", + "ma", + "them", + "mightn't", + "shan", + "mustn", + "they", + "what", + "both", + "that'll", + "how", + "is", + "he", + "because", + "down", + "haven't", + "are", + "no", + "it's", + "our", + "being", + "the", + "or", + "above", + "myself", + "once", + "don't", + "doesn't", + "as", + "nor", + "here", + "herself", + "hasn", + "mightn", + "have", + "its", + "all", + "were", + "ain", + "this", + "at", + "after", + "over", + "shouldn't", + "into", + "before", + "don", + "wouldn", + "re", + "couldn't", + "wasn", + "in", + "should", + "there", + "himself", + "isn't", + "should've", + "doing", + "ve", + "shouldn", + "a", + "did", + "and", + "his", + "between", + "me", + "up", + "below", + "人民", + "末##末", + "啊", + "阿", + "哎", + "哎呀", + "哎哟", + "唉", + "俺", + "俺们", + "按", + "按照", + "吧", + "吧哒", + "把", + "罢了", + "被", + "本", + "本着", + "比", + "比方", + "比如", + "鄙人", + "彼", + "彼此", + "边", + "别", + "别的", + "别说", + "并", + "并且", + "不比", + "不成", + "不单", + "不但", + "不独", + "不管", + "不光", + "不过", + "不仅", + "不拘", + "不论", + "不怕", + "不然", + "不如", + "不特", + "不惟", + "不问", + "不只", + "朝", + "朝着", + "趁", + "趁着", + "乘", + "冲", + "除", + "除此之外", + "除非", + "除了", + "此", + "此间", + "此外", + "从", + "从而", + "打", + "待", + "但", + "但是", + "当", + "当着", + "到", + "得", + "的", + "的话", + "等", + "等等", + "地", + "第", + "叮咚", + "对", + "对于", + "多", + "多少", + "而", + "而况", + "而且", + "而是", + "而外", + "而言", + "而已", + "尔后", + "反过来", + "反过来说", + "反之", + "非但", + "非徒", + "否则", + "嘎", + "嘎登", + "该", + "赶", + "个", + "各", + "各个", + "各位", + "各种", + "各自", + "给", + "根据", + "跟", + "故", + "故此", + "固然", + "关于", + "管", + "归", + "果然", + "果真", + "过", + "哈", + "哈哈", + "呵", + "和", + "何", + "何处", + "何况", + "何时", + "嘿", + "哼", + "哼唷", + "呼哧", + "乎", + "哗", + "还是", + "还有", + "换句话说", + "换言之", + "或", + "或是", + "或者", + "极了", + "及", + "及其", + "及至", + "即", + "即便", + "即或", + "即令", + "即若", + "即使", + "几", + "几时", + "己", + "既", + "既然", + "既是", + "继而", + "加之", + "假如", + "假若", + "假使", + "鉴于", + "将", + "较", + "较之", + "叫", + "接着", + "结果", + "借", + "紧接着", + "进而", + "尽", + "尽管", + "经", + "经过", + "就", + "就是", + "就是说", + "据", + "具体地说", + "具体说来", + "开始", + "开外", + "靠", + "咳", + "可", + "可见", + "可是", + "可以", + "况且", + "啦", + "来", + "来着", + "离", + "例如", + "哩", + "连", + "连同", + "两者", + "了", + "临", + "另", + "另外", + "另一方面", + "论", + "嘛", + "吗", + "慢说", + "漫说", + "冒", + "么", + "每", + "每当", + "们", + "莫若", + "某", + "某个", + "某些", + "拿", + "哪", + "哪边", + "哪儿", + "哪个", + "哪里", + "哪年", + "哪怕", + "哪天", + "哪些", + "哪样", + "那", + "那边", + "那儿", + "那个", + "那会儿", + "那里", + "那么", + "那么些", + "那么样", + "那时", + "那些", + "那样", + "乃", + "乃至", + "呢", + "能", + "你", + "你们", + "您", + "宁", + "宁可", + "宁肯", + "宁愿", + "哦", + "呕", + "啪达", + "旁人", + "呸", + "凭", + "凭借", + "其", + "其次", + "其二", + "其他", + "其它", + "其一", + "其余", + "其中", + "起", + "起见", + "岂但", + "恰恰相反", + "前后", + "前者", + "且", + "然而", + "然后", + "然则", + "让", + "人家", + "任", + "任何", + "任凭", + "如", + "如此", + "如果", + "如何", + "如其", + "如若", + "如上所述", + "若", + "若非", + "若是", + "啥", + "上下", + "尚且", + "设若", + "设使", + "甚而", + "甚么", + "甚至", + "省得", + "时候", + "什么", + "什么样", + "使得", + "是", + "是的", + "首先", + "谁", + "谁知", + "顺", + "顺着", + "似的", + "虽", + "虽然", + "虽说", + "虽则", + "随", + "随着", + "所", + "所以", + "他", + "他们", + "他人", + "它", + "它们", + "她", + "她们", + "倘", + "倘或", + "倘然", + "倘若", + "倘使", + "腾", + "替", + "通过", + "同", + "同时", + "哇", + "万一", + "往", + "望", + "为", + "为何", + "为了", + "为什么", + "为着", + "喂", + "嗡嗡", + "我", + "我们", + "呜", + "呜呼", + "乌乎", + "无论", + "无宁", + "毋宁", + "嘻", + "吓", + "相对而言", + "像", + "向", + "向着", + "嘘", + "呀", + "焉", + "沿", + "沿着", + "要", + "要不", + "要不然", + "要不是", + "要么", + "要是", + "也", + "也罢", + "也好", + "一", + "一般", + "一旦", + "一方面", + "一来", + "一切", + "一样", + "一则", + "依", + "依照", + "矣", + "以", + "以便", + "以及", + "以免", + "以至", + "以至于", + "以致", + "抑或", + "因", + "因此", + "因而", + "因为", + "哟", + "用", + "由", + "由此可见", + "由于", + "有", + "有的", + "有关", + "有些", + "又", + "于", + "于是", + "于是乎", + "与", + "与此同时", + "与否", + "与其", + "越是", + "云云", + "哉", + "再说", + "再者", + "在", + "在下", + "咱", + "咱们", + "则", + "怎", + "怎么", + "怎么办", + "怎么样", + "怎样", + "咋", + "照", + "照着", + "者", + "这", + "这边", + "这儿", + "这个", + "这会儿", + "这就是说", + "这里", + "这么", + "这么点儿", + "这么些", + "这么样", + "这时", + "这些", + "这样", + "正如", + "吱", + "之", + "之类", + "之所以", + "之一", + "只是", + "只限", + "只要", + "只有", + "至", + "至于", + "诸位", + "着", + "着呢", + "自", + "自从", + "自个儿", + "自各儿", + "自己", + "自家", + "自身", + "综上所述", + "总的来看", + "总的来说", + "总的说来", + "总而言之", + "总之", + "纵", + "纵令", + "纵然", + "纵使", + "遵照", + "作为", + "兮", + "呃", + "呗", + "咚", + "咦", + "喏", + "啐", + "喔唷", + "嗬", + "嗯", + "嗳", + "~", + "!", + ".", + ":", + '"', + "'", + "(", + ")", + "*", + "A", + "白", + "社会主义", + "--", + "..", + ">>", + " [", + " ]", + "", + "<", + ">", + "/", + "\\", + "|", + "-", + "_", + "+", + "=", + "&", + "^", + "%", + "#", + "@", + "`", + ";", + "$", + "(", + ")", + "——", + "—", + "¥", + "·", + "...", + "‘", + "’", + "〉", + "〈", + "…", + " ", + "0", + "1", + "2", + "3", + "4", + "5", + "6", + "7", + "8", + "9", + "0", + "1", + "2", + "3", + "4", + "5", + "6", + "7", + "8", + "9", + "二", + "三", + "四", + "五", + "六", + "七", + "八", + "九", + "零", + ">", + "<", + "@", + "#", + "$", + "%", + "︿", + "&", + "*", + "+", + "~", + "|", + "[", + "]", + "{", + "}", + "啊哈", + "啊呀", + "啊哟", + "挨次", + "挨个", + "挨家挨户", + "挨门挨户", + "挨门逐户", + "挨着", + "按理", + "按期", + "按时", + "按说", + "暗地里", + "暗中", + "暗自", + "昂然", + "八成", + "白白", + "半", + "梆", + "保管", + "保险", + "饱", + "背地里", + "背靠背", + "倍感", + "倍加", + "本人", + "本身", + "甭", + "比起", + "比如说", + "比照", + "毕竟", + "必", + "必定", + "必将", + "必须", + "便", + "别人", + "并非", + "并肩", + "并没", + "并没有", + "并排", + "并无", + "勃然", + "不", + "不必", + "不常", + "不大", + "不但...而且", + "不得", + "不得不", + "不得了", + "不得已", + "不迭", + "不定", + "不对", + "不妨", + "不管怎样", + "不会", + "不仅...而且", + "不仅仅", + "不仅仅是", + "不经意", + "不可开交", + "不可抗拒", + "不力", + "不了", + "不料", + "不满", + "不免", + "不能不", + "不起", + "不巧", + "不然的话", + "不日", + "不少", + "不胜", + "不时", + "不是", + "不同", + "不能", + "不要", + "不外", + "不外乎", + "不下", + "不限", + "不消", + "不已", + "不亦乐乎", + "不由得", + "不再", + "不择手段", + "不怎么", + "不曾", + "不知不觉", + "不止", + "不止一次", + "不至于", + "才", + "才能", + "策略地", + "差不多", + "差一点", + "常", + "常常", + "常言道", + "常言说", + "常言说得好", + "长此下去", + "长话短说", + "长期以来", + "长线", + "敞开儿", + "彻夜", + "陈年", + "趁便", + "趁机", + "趁热", + "趁势", + "趁早", + "成年", + "成年累月", + "成心", + "乘机", + "乘胜", + "乘势", + "乘隙", + "乘虚", + "诚然", + "迟早", + "充分", + "充其极", + "充其量", + "抽冷子", + "臭", + "初", + "出", + "出来", + "出去", + "除此", + "除此而外", + "除此以外", + "除开", + "除去", + "除却", + "除外", + "处处", + "川流不息", + "传", + "传说", + "传闻", + "串行", + "纯", + "纯粹", + "此后", + "此中", + "次第", + "匆匆", + "从不", + "从此", + "从此以后", + "从古到今", + "从古至今", + "从今以后", + "从宽", + "从来", + "从轻", + "从速", + "从头", + "从未", + "从无到有", + "从小", + "从新", + "从严", + "从优", + "从早到晚", + "从中", + "从重", + "凑巧", + "粗", + "存心", + "达旦", + "打从", + "打开天窗说亮话", + "大", + "大不了", + "大大", + "大抵", + "大都", + "大多", + "大凡", + "大概", + "大家", + "大举", + "大略", + "大面儿上", + "大事", + "大体", + "大体上", + "大约", + "大张旗鼓", + "大致", + "呆呆地", + "带", + "殆", + "待到", + "单", + "单纯", + "单单", + "但愿", + "弹指之间", + "当场", + "当儿", + "当即", + "当口儿", + "当然", + "当庭", + "当头", + "当下", + "当真", + "当中", + "倒不如", + "倒不如说", + "倒是", + "到处", + "到底", + "到了儿", + "到目前为止", + "到头", + "到头来", + "得起", + "得天独厚", + "的确", + "等到", + "叮当", + "顶多", + "定", + "动不动", + "动辄", + "陡然", + "都", + "独", + "独自", + "断然", + "顿时", + "多次", + "多多", + "多多少少", + "多多益善", + "多亏", + "多年来", + "多年前", + "而后", + "而论", + "而又", + "尔等", + "二话不说", + "二话没说", + "反倒", + "反倒是", + "反而", + "反手", + "反之亦然", + "反之则", + "方", + "方才", + "方能", + "放量", + "非常", + "非得", + "分期", + "分期分批", + "分头", + "奋勇", + "愤然", + "风雨无阻", + "逢", + "弗", + "甫", + "嘎嘎", + "该当", + "概", + "赶快", + "赶早不赶晚", + "敢", + "敢情", + "敢于", + "刚", + "刚才", + "刚好", + "刚巧", + "高低", + "格外", + "隔日", + "隔夜", + "个人", + "各式", + "更", + "更加", + "更进一步", + "更为", + "公然", + "共", + "共总", + "够瞧的", + "姑且", + "古来", + "故而", + "故意", + "固", + "怪", + "怪不得", + "惯常", + "光", + "光是", + "归根到底", + "归根结底", + "过于", + "毫不", + "毫无", + "毫无保留地", + "毫无例外", + "好在", + "何必", + "何尝", + "何妨", + "何苦", + "何乐而不为", + "何须", + "何止", + "很", + "很多", + "很少", + "轰然", + "后来", + "呼啦", + "忽地", + "忽然", + "互", + "互相", + "哗啦", + "话说", + "还", + "恍然", + "会", + "豁然", + "活", + "伙同", + "或多或少", + "或许", + "基本", + "基本上", + "基于", + "极", + "极大", + "极度", + "极端", + "极力", + "极其", + "极为", + "急匆匆", + "即将", + "即刻", + "即是说", + "几度", + "几番", + "几乎", + "几经", + "既...又", + "继之", + "加上", + "加以", + "间或", + "简而言之", + "简言之", + "简直", + "见", + "将才", + "将近", + "将要", + "交口", + "较比", + "较为", + "接连不断", + "接下来", + "皆可", + "截然", + "截至", + "藉以", + "借此", + "借以", + "届时", + "仅", + "仅仅", + "谨", + "进来", + "进去", + "近", + "近几年来", + "近来", + "近年来", + "尽管如此", + "尽可能", + "尽快", + "尽量", + "尽然", + "尽如人意", + "尽心竭力", + "尽心尽力", + "尽早", + "精光", + "经常", + "竟", + "竟然", + "究竟", + "就此", + "就地", + "就算", + "居然", + "局外", + "举凡", + "据称", + "据此", + "据实", + "据说", + "据我所知", + "据悉", + "具体来说", + "决不", + "决非", + "绝", + "绝不", + "绝顶", + "绝对", + "绝非", + "均", + "喀", + "看", + "看来", + "看起来", + "看上去", + "看样子", + "可好", + "可能", + "恐怕", + "快", + "快要", + "来不及", + "来得及", + "来讲", + "来看", + "拦腰", + "牢牢", + "老", + "老大", + "老老实实", + "老是", + "累次", + "累年", + "理当", + "理该", + "理应", + "历", + "立", + "立地", + "立刻", + "立马", + "立时", + "联袂", + "连连", + "连日", + "连日来", + "连声", + "连袂", + "临到", + "另方面", + "另行", + "另一个", + "路经", + "屡", + "屡次", + "屡次三番", + "屡屡", + "缕缕", + "率尔", + "率然", + "略", + "略加", + "略微", + "略为", + "论说", + "马上", + "蛮", + "满", + "没", + "没有", + "每逢", + "每每", + "每时每刻", + "猛然", + "猛然间", + "莫", + "莫不", + "莫非", + "莫如", + "默默地", + "默然", + "呐", + "那末", + "奈", + "难道", + "难得", + "难怪", + "难说", + "内", + "年复一年", + "凝神", + "偶而", + "偶尔", + "怕", + "砰", + "碰巧", + "譬如", + "偏偏", + "乒", + "平素", + "颇", + "迫于", + "扑通", + "其后", + "其实", + "奇", + "齐", + "起初", + "起来", + "起首", + "起头", + "起先", + "岂", + "岂非", + "岂止", + "迄", + "恰逢", + "恰好", + "恰恰", + "恰巧", + "恰如", + "恰似", + "千", + "千万", + "千万千万", + "切", + "切不可", + "切莫", + "切切", + "切勿", + "窃", + "亲口", + "亲身", + "亲手", + "亲眼", + "亲自", + "顷", + "顷刻", + "顷刻间", + "顷刻之间", + "请勿", + "穷年累月", + "取道", + "去", + "权时", + "全都", + "全力", + "全年", + "全然", + "全身心", + "然", + "人人", + "仍", + "仍旧", + "仍然", + "日复一日", + "日见", + "日渐", + "日益", + "日臻", + "如常", + "如此等等", + "如次", + "如今", + "如期", + "如前所述", + "如上", + "如下", + "汝", + "三番两次", + "三番五次", + "三天两头", + "瑟瑟", + "沙沙", + "上", + "上来", + "上去", + "一个", + "月", + "日", + "\n", +} diff --git a/api/core/rag/datasource/keyword/keyword_base.py b/api/core/rag/datasource/keyword/keyword_base.py new file mode 100644 index 0000000000000000000000000000000000000000..b261b40b7286929b1838d1d8f236ef910fa8b08e --- /dev/null +++ b/api/core/rag/datasource/keyword/keyword_base.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + +from core.rag.models.document import Document +from models.dataset import Dataset + + +class BaseKeyword(ABC): + def __init__(self, dataset: Dataset): + self.dataset = dataset + + @abstractmethod + def create(self, texts: list[Document], **kwargs) -> BaseKeyword: + raise NotImplementedError + + @abstractmethod + def add_texts(self, texts: list[Document], **kwargs): + raise NotImplementedError + + @abstractmethod + def text_exists(self, id: str) -> bool: + raise NotImplementedError + + @abstractmethod + def delete_by_ids(self, ids: list[str]) -> None: + raise NotImplementedError + + @abstractmethod + def delete(self) -> None: + raise NotImplementedError + + @abstractmethod + def search(self, query: str, **kwargs: Any) -> list[Document]: + raise NotImplementedError + + def _filter_duplicate_texts(self, texts: list[Document]) -> list[Document]: + for text in texts.copy(): + if text.metadata is None: + continue + doc_id = text.metadata["doc_id"] + exists_duplicate_node = self.text_exists(doc_id) + if exists_duplicate_node: + texts.remove(text) + + return texts + + def _get_uuids(self, texts: list[Document]) -> list[str]: + return [text.metadata["doc_id"] for text in texts if text.metadata] diff --git a/api/core/rag/datasource/keyword/keyword_factory.py b/api/core/rag/datasource/keyword/keyword_factory.py new file mode 100644 index 0000000000000000000000000000000000000000..f1a6ade91f9bd169a01d5b43225cf66a49407df1 --- /dev/null +++ b/api/core/rag/datasource/keyword/keyword_factory.py @@ -0,0 +1,54 @@ +from typing import Any + +from configs import dify_config +from core.rag.datasource.keyword.keyword_base import BaseKeyword +from core.rag.datasource.keyword.keyword_type import KeyWordType +from core.rag.models.document import Document +from models.dataset import Dataset + + +class Keyword: + def __init__(self, dataset: Dataset): + self._dataset = dataset + self._keyword_processor = self._init_keyword() + + def _init_keyword(self) -> BaseKeyword: + keyword_type = dify_config.KEYWORD_STORE + keyword_factory = self.get_keyword_factory(keyword_type) + return keyword_factory(self._dataset) + + @staticmethod + def get_keyword_factory(keyword_type: str) -> type[BaseKeyword]: + match keyword_type: + case KeyWordType.JIEBA: + from core.rag.datasource.keyword.jieba.jieba import Jieba + + return Jieba + case _: + raise ValueError(f"Keyword store {keyword_type} is not supported.") + + def create(self, texts: list[Document], **kwargs): + self._keyword_processor.create(texts, **kwargs) + + def add_texts(self, texts: list[Document], **kwargs): + self._keyword_processor.add_texts(texts, **kwargs) + + def text_exists(self, id: str) -> bool: + return self._keyword_processor.text_exists(id) + + def delete_by_ids(self, ids: list[str]) -> None: + self._keyword_processor.delete_by_ids(ids) + + def delete(self) -> None: + self._keyword_processor.delete() + + def search(self, query: str, **kwargs: Any) -> list[Document]: + return self._keyword_processor.search(query, **kwargs) + + def __getattr__(self, name): + if self._keyword_processor is not None: + method = getattr(self._keyword_processor, name) + if callable(method): + return method + + raise AttributeError(f"'Keyword' object has no attribute '{name}'") diff --git a/api/core/rag/datasource/keyword/keyword_type.py b/api/core/rag/datasource/keyword/keyword_type.py new file mode 100644 index 0000000000000000000000000000000000000000..d845c7111dd5789b3971a40cb74d687948c6df76 --- /dev/null +++ b/api/core/rag/datasource/keyword/keyword_type.py @@ -0,0 +1,5 @@ +from enum import StrEnum + + +class KeyWordType(StrEnum): + JIEBA = "jieba" diff --git a/api/core/rag/datasource/retrieval_service.py b/api/core/rag/datasource/retrieval_service.py new file mode 100644 index 0000000000000000000000000000000000000000..3a8200bc7b565018ba1d49ffe19a154a4898a115 --- /dev/null +++ b/api/core/rag/datasource/retrieval_service.py @@ -0,0 +1,339 @@ +import threading +from typing import Optional + +from flask import Flask, current_app + +from core.rag.data_post_processor.data_post_processor import DataPostProcessor +from core.rag.datasource.keyword.keyword_factory import Keyword +from core.rag.datasource.vdb.vector_factory import Vector +from core.rag.embedding.retrieval import RetrievalSegments +from core.rag.index_processor.constant.index_type import IndexType +from core.rag.models.document import Document +from core.rag.rerank.rerank_type import RerankMode +from core.rag.retrieval.retrieval_methods import RetrievalMethod +from extensions.ext_database import db +from models.dataset import ChildChunk, Dataset, DocumentSegment +from models.dataset import Document as DatasetDocument +from services.external_knowledge_service import ExternalDatasetService + +default_retrieval_model = { + "search_method": RetrievalMethod.SEMANTIC_SEARCH.value, + "reranking_enable": False, + "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""}, + "top_k": 2, + "score_threshold_enabled": False, +} + + +class RetrievalService: + @classmethod + def retrieve( + cls, + retrieval_method: str, + dataset_id: str, + query: str, + top_k: int, + score_threshold: Optional[float] = 0.0, + reranking_model: Optional[dict] = None, + reranking_mode: str = "reranking_model", + weights: Optional[dict] = None, + ): + if not query: + return [] + dataset = db.session.query(Dataset).filter(Dataset.id == dataset_id).first() + if not dataset: + return [] + + if not dataset or dataset.available_document_count == 0 or dataset.available_segment_count == 0: + return [] + all_documents: list[Document] = [] + threads: list[threading.Thread] = [] + exceptions: list[str] = [] + # retrieval_model source with keyword + if retrieval_method == "keyword_search": + keyword_thread = threading.Thread( + target=RetrievalService.keyword_search, + kwargs={ + "flask_app": current_app._get_current_object(), # type: ignore + "dataset_id": dataset_id, + "query": query, + "top_k": top_k, + "all_documents": all_documents, + "exceptions": exceptions, + }, + ) + threads.append(keyword_thread) + keyword_thread.start() + # retrieval_model source with semantic + if RetrievalMethod.is_support_semantic_search(retrieval_method): + embedding_thread = threading.Thread( + target=RetrievalService.embedding_search, + kwargs={ + "flask_app": current_app._get_current_object(), # type: ignore + "dataset_id": dataset_id, + "query": query, + "top_k": top_k, + "score_threshold": score_threshold, + "reranking_model": reranking_model, + "all_documents": all_documents, + "retrieval_method": retrieval_method, + "exceptions": exceptions, + }, + ) + threads.append(embedding_thread) + embedding_thread.start() + + # retrieval source with full text + if RetrievalMethod.is_support_fulltext_search(retrieval_method): + full_text_index_thread = threading.Thread( + target=RetrievalService.full_text_index_search, + kwargs={ + "flask_app": current_app._get_current_object(), # type: ignore + "dataset_id": dataset_id, + "query": query, + "retrieval_method": retrieval_method, + "score_threshold": score_threshold, + "top_k": top_k, + "reranking_model": reranking_model, + "all_documents": all_documents, + "exceptions": exceptions, + }, + ) + threads.append(full_text_index_thread) + full_text_index_thread.start() + + for thread in threads: + thread.join() + + if exceptions: + exception_message = ";\n".join(exceptions) + raise ValueError(exception_message) + + if retrieval_method == RetrievalMethod.HYBRID_SEARCH.value: + data_post_processor = DataPostProcessor( + str(dataset.tenant_id), reranking_mode, reranking_model, weights, False + ) + all_documents = data_post_processor.invoke( + query=query, + documents=all_documents, + score_threshold=score_threshold, + top_n=top_k, + ) + + return all_documents + + @classmethod + def external_retrieve(cls, dataset_id: str, query: str, external_retrieval_model: Optional[dict] = None): + dataset = db.session.query(Dataset).filter(Dataset.id == dataset_id).first() + if not dataset: + return [] + all_documents = ExternalDatasetService.fetch_external_knowledge_retrieval( + dataset.tenant_id, dataset_id, query, external_retrieval_model or {} + ) + return all_documents + + @classmethod + def keyword_search( + cls, flask_app: Flask, dataset_id: str, query: str, top_k: int, all_documents: list, exceptions: list + ): + with flask_app.app_context(): + try: + dataset = db.session.query(Dataset).filter(Dataset.id == dataset_id).first() + if not dataset: + raise ValueError("dataset not found") + + keyword = Keyword(dataset=dataset) + + documents = keyword.search(cls.escape_query_for_search(query), top_k=top_k) + all_documents.extend(documents) + except Exception as e: + exceptions.append(str(e)) + + @classmethod + def embedding_search( + cls, + flask_app: Flask, + dataset_id: str, + query: str, + top_k: int, + score_threshold: Optional[float], + reranking_model: Optional[dict], + all_documents: list, + retrieval_method: str, + exceptions: list, + ): + with flask_app.app_context(): + try: + dataset = db.session.query(Dataset).filter(Dataset.id == dataset_id).first() + if not dataset: + raise ValueError("dataset not found") + + vector = Vector(dataset=dataset) + + documents = vector.search_by_vector( + cls.escape_query_for_search(query), + search_type="similarity_score_threshold", + top_k=top_k, + score_threshold=score_threshold, + filter={"group_id": [dataset.id]}, + ) + + if documents: + if ( + reranking_model + and reranking_model.get("reranking_model_name") + and reranking_model.get("reranking_provider_name") + and retrieval_method == RetrievalMethod.SEMANTIC_SEARCH.value + ): + data_post_processor = DataPostProcessor( + str(dataset.tenant_id), RerankMode.RERANKING_MODEL.value, reranking_model, None, False + ) + all_documents.extend( + data_post_processor.invoke( + query=query, + documents=documents, + score_threshold=score_threshold, + top_n=len(documents), + ) + ) + else: + all_documents.extend(documents) + except Exception as e: + exceptions.append(str(e)) + + @classmethod + def full_text_index_search( + cls, + flask_app: Flask, + dataset_id: str, + query: str, + top_k: int, + score_threshold: Optional[float], + reranking_model: Optional[dict], + all_documents: list, + retrieval_method: str, + exceptions: list, + ): + with flask_app.app_context(): + try: + dataset = db.session.query(Dataset).filter(Dataset.id == dataset_id).first() + if not dataset: + raise ValueError("dataset not found") + + vector_processor = Vector( + dataset=dataset, + ) + + documents = vector_processor.search_by_full_text(cls.escape_query_for_search(query), top_k=top_k) + if documents: + if ( + reranking_model + and reranking_model.get("reranking_model_name") + and reranking_model.get("reranking_provider_name") + and retrieval_method == RetrievalMethod.FULL_TEXT_SEARCH.value + ): + data_post_processor = DataPostProcessor( + str(dataset.tenant_id), RerankMode.RERANKING_MODEL.value, reranking_model, None, False + ) + all_documents.extend( + data_post_processor.invoke( + query=query, + documents=documents, + score_threshold=score_threshold, + top_n=len(documents), + ) + ) + else: + all_documents.extend(documents) + except Exception as e: + exceptions.append(str(e)) + + @staticmethod + def escape_query_for_search(query: str) -> str: + return query.replace('"', '\\"') + + @staticmethod + def format_retrieval_documents(documents: list[Document]) -> list[RetrievalSegments]: + records = [] + include_segment_ids = [] + segment_child_map = {} + for document in documents: + document_id = document.metadata.get("document_id") + dataset_document = db.session.query(DatasetDocument).filter(DatasetDocument.id == document_id).first() + if dataset_document: + if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX: + child_index_node_id = document.metadata.get("doc_id") + result = ( + db.session.query(ChildChunk, DocumentSegment) + .join(DocumentSegment, ChildChunk.segment_id == DocumentSegment.id) + .filter( + ChildChunk.index_node_id == child_index_node_id, + DocumentSegment.dataset_id == dataset_document.dataset_id, + DocumentSegment.enabled == True, + DocumentSegment.status == "completed", + ) + .first() + ) + if result: + child_chunk, segment = result + if not segment: + continue + if segment.id not in include_segment_ids: + include_segment_ids.append(segment.id) + child_chunk_detail = { + "id": child_chunk.id, + "content": child_chunk.content, + "position": child_chunk.position, + "score": document.metadata.get("score", 0.0), + } + map_detail = { + "max_score": document.metadata.get("score", 0.0), + "child_chunks": [child_chunk_detail], + } + segment_child_map[segment.id] = map_detail + record = { + "segment": segment, + } + records.append(record) + else: + child_chunk_detail = { + "id": child_chunk.id, + "content": child_chunk.content, + "position": child_chunk.position, + "score": document.metadata.get("score", 0.0), + } + segment_child_map[segment.id]["child_chunks"].append(child_chunk_detail) + segment_child_map[segment.id]["max_score"] = max( + segment_child_map[segment.id]["max_score"], document.metadata.get("score", 0.0) + ) + else: + continue + else: + index_node_id = document.metadata["doc_id"] + + segment = ( + db.session.query(DocumentSegment) + .filter( + DocumentSegment.dataset_id == dataset_document.dataset_id, + DocumentSegment.enabled == True, + DocumentSegment.status == "completed", + DocumentSegment.index_node_id == index_node_id, + ) + .first() + ) + + if not segment: + continue + include_segment_ids.append(segment.id) + record = { + "segment": segment, + "score": document.metadata.get("score", None), + } + + records.append(record) + for record in records: + if record["segment"].id in segment_child_map: + record["child_chunks"] = segment_child_map[record["segment"].id].get("child_chunks", None) + record["score"] = segment_child_map[record["segment"].id]["max_score"] + + return [RetrievalSegments(**record) for record in records] diff --git a/api/core/rag/datasource/vdb/__init__.py b/api/core/rag/datasource/vdb/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/analyticdb/__init__.py b/api/core/rag/datasource/vdb/analyticdb/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/analyticdb/analyticdb_vector.py b/api/core/rag/datasource/vdb/analyticdb/analyticdb_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..603d3fdbcdf1abc340d20c349e81ac2070e9f333 --- /dev/null +++ b/api/core/rag/datasource/vdb/analyticdb/analyticdb_vector.py @@ -0,0 +1,104 @@ +import json +from typing import Any + +from configs import dify_config +from core.rag.datasource.vdb.analyticdb.analyticdb_vector_openapi import ( + AnalyticdbVectorOpenAPI, + AnalyticdbVectorOpenAPIConfig, +) +from core.rag.datasource.vdb.analyticdb.analyticdb_vector_sql import AnalyticdbVectorBySql, AnalyticdbVectorBySqlConfig +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from models.dataset import Dataset + + +class AnalyticdbVector(BaseVector): + def __init__( + self, + collection_name: str, + api_config: AnalyticdbVectorOpenAPIConfig | None, + sql_config: AnalyticdbVectorBySqlConfig | None, + ): + super().__init__(collection_name) + if api_config is not None: + self.analyticdb_vector: AnalyticdbVectorOpenAPI | AnalyticdbVectorBySql = AnalyticdbVectorOpenAPI( + collection_name, api_config + ) + else: + if sql_config is None: + raise ValueError("Either api_config or sql_config must be provided") + self.analyticdb_vector = AnalyticdbVectorBySql(collection_name, sql_config) + + def get_type(self) -> str: + return VectorType.ANALYTICDB + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + dimension = len(embeddings[0]) + self.analyticdb_vector._create_collection_if_not_exists(dimension) + self.analyticdb_vector.add_texts(texts, embeddings) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + self.analyticdb_vector.add_texts(documents, embeddings) + + def text_exists(self, id: str) -> bool: + return self.analyticdb_vector.text_exists(id) + + def delete_by_ids(self, ids: list[str]) -> None: + self.analyticdb_vector.delete_by_ids(ids) + + def delete_by_metadata_field(self, key: str, value: str) -> None: + self.analyticdb_vector.delete_by_metadata_field(key, value) + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + return self.analyticdb_vector.search_by_vector(query_vector) + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + return self.analyticdb_vector.search_by_full_text(query, **kwargs) + + def delete(self) -> None: + self.analyticdb_vector.delete() + + +class AnalyticdbVectorFactory(AbstractVectorFactory): + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> AnalyticdbVector: + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix.lower() + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id).lower() + dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.ANALYTICDB, collection_name)) + + if dify_config.ANALYTICDB_HOST is None: + # implemented through OpenAPI + apiConfig = AnalyticdbVectorOpenAPIConfig( + access_key_id=dify_config.ANALYTICDB_KEY_ID or "", + access_key_secret=dify_config.ANALYTICDB_KEY_SECRET or "", + region_id=dify_config.ANALYTICDB_REGION_ID or "", + instance_id=dify_config.ANALYTICDB_INSTANCE_ID or "", + account=dify_config.ANALYTICDB_ACCOUNT or "", + account_password=dify_config.ANALYTICDB_PASSWORD or "", + namespace=dify_config.ANALYTICDB_NAMESPACE or "", + namespace_password=dify_config.ANALYTICDB_NAMESPACE_PASSWORD, + ) + sqlConfig = None + else: + # implemented through sql + sqlConfig = AnalyticdbVectorBySqlConfig( + host=dify_config.ANALYTICDB_HOST, + port=dify_config.ANALYTICDB_PORT, + account=dify_config.ANALYTICDB_ACCOUNT or "", + account_password=dify_config.ANALYTICDB_PASSWORD or "", + min_connection=dify_config.ANALYTICDB_MIN_CONNECTION, + max_connection=dify_config.ANALYTICDB_MAX_CONNECTION, + namespace=dify_config.ANALYTICDB_NAMESPACE or "", + ) + apiConfig = None + return AnalyticdbVector( + collection_name, + apiConfig, + sqlConfig, + ) diff --git a/api/core/rag/datasource/vdb/analyticdb/analyticdb_vector_openapi.py b/api/core/rag/datasource/vdb/analyticdb/analyticdb_vector_openapi.py new file mode 100644 index 0000000000000000000000000000000000000000..095752ea8eaa4227f2c476563d1e71b6af975561 --- /dev/null +++ b/api/core/rag/datasource/vdb/analyticdb/analyticdb_vector_openapi.py @@ -0,0 +1,310 @@ +import json +from typing import Any, Optional + +from pydantic import BaseModel, model_validator + +_import_err_msg = ( + "`alibabacloud_gpdb20160503` and `alibabacloud_tea_openapi` packages not found, " + "please run `pip install alibabacloud_gpdb20160503 alibabacloud_tea_openapi`" +) + +from core.rag.models.document import Document +from extensions.ext_redis import redis_client + + +class AnalyticdbVectorOpenAPIConfig(BaseModel): + access_key_id: str + access_key_secret: str + region_id: str + instance_id: str + account: str + account_password: str + namespace: str = "dify" + namespace_password: Optional[str] = None + metrics: str = "cosine" + read_timeout: int = 60000 + + @model_validator(mode="before") + @classmethod + def validate_config(cls, values: dict) -> dict: + if not values["access_key_id"]: + raise ValueError("config ANALYTICDB_KEY_ID is required") + if not values["access_key_secret"]: + raise ValueError("config ANALYTICDB_KEY_SECRET is required") + if not values["region_id"]: + raise ValueError("config ANALYTICDB_REGION_ID is required") + if not values["instance_id"]: + raise ValueError("config ANALYTICDB_INSTANCE_ID is required") + if not values["account"]: + raise ValueError("config ANALYTICDB_ACCOUNT is required") + if not values["account_password"]: + raise ValueError("config ANALYTICDB_PASSWORD is required") + if not values["namespace_password"]: + raise ValueError("config ANALYTICDB_NAMESPACE_PASSWORD is required") + return values + + def to_analyticdb_client_params(self): + return { + "access_key_id": self.access_key_id, + "access_key_secret": self.access_key_secret, + "region_id": self.region_id, + "read_timeout": self.read_timeout, + } + + +class AnalyticdbVectorOpenAPI: + def __init__(self, collection_name: str, config: AnalyticdbVectorOpenAPIConfig): + try: + from alibabacloud_gpdb20160503.client import Client # type: ignore + from alibabacloud_tea_openapi import models as open_api_models # type: ignore + except: + raise ImportError(_import_err_msg) + self._collection_name = collection_name.lower() + self.config = config + self._client_config = open_api_models.Config(user_agent="dify", **config.to_analyticdb_client_params()) + self._client = Client(self._client_config) + self._initialize() + + def _initialize(self) -> None: + cache_key = f"vector_initialize_{self.config.instance_id}" + lock_name = f"{cache_key}_lock" + with redis_client.lock(lock_name, timeout=20): + database_exist_cache_key = f"vector_initialize_{self.config.instance_id}" + if redis_client.get(database_exist_cache_key): + return + self._initialize_vector_database() + self._create_namespace_if_not_exists() + redis_client.set(database_exist_cache_key, 1, ex=3600) + + def _initialize_vector_database(self) -> None: + from alibabacloud_gpdb20160503 import models as gpdb_20160503_models # type: ignore + + request = gpdb_20160503_models.InitVectorDatabaseRequest( + dbinstance_id=self.config.instance_id, + region_id=self.config.region_id, + manager_account=self.config.account, + manager_account_password=self.config.account_password, + ) + self._client.init_vector_database(request) + + def _create_namespace_if_not_exists(self) -> None: + from alibabacloud_gpdb20160503 import models as gpdb_20160503_models + from Tea.exceptions import TeaException # type: ignore + + try: + request = gpdb_20160503_models.DescribeNamespaceRequest( + dbinstance_id=self.config.instance_id, + region_id=self.config.region_id, + namespace=self.config.namespace, + manager_account=self.config.account, + manager_account_password=self.config.account_password, + ) + self._client.describe_namespace(request) + except TeaException as e: + if e.statusCode == 404: + request = gpdb_20160503_models.CreateNamespaceRequest( + dbinstance_id=self.config.instance_id, + region_id=self.config.region_id, + manager_account=self.config.account, + manager_account_password=self.config.account_password, + namespace=self.config.namespace, + namespace_password=self.config.namespace_password, + ) + self._client.create_namespace(request) + else: + raise ValueError(f"failed to create namespace {self.config.namespace}: {e}") + + def _create_collection_if_not_exists(self, embedding_dimension: int): + from alibabacloud_gpdb20160503 import models as gpdb_20160503_models + from Tea.exceptions import TeaException + + cache_key = f"vector_indexing_{self._collection_name}" + lock_name = f"{cache_key}_lock" + with redis_client.lock(lock_name, timeout=20): + collection_exist_cache_key = f"vector_indexing_{self._collection_name}" + if redis_client.get(collection_exist_cache_key): + return + try: + request = gpdb_20160503_models.DescribeCollectionRequest( + dbinstance_id=self.config.instance_id, + region_id=self.config.region_id, + namespace=self.config.namespace, + namespace_password=self.config.namespace_password, + collection=self._collection_name, + ) + self._client.describe_collection(request) + except TeaException as e: + if e.statusCode == 404: + metadata = '{"ref_doc_id":"text","page_content":"text","metadata_":"jsonb"}' + full_text_retrieval_fields = "page_content" + request = gpdb_20160503_models.CreateCollectionRequest( + dbinstance_id=self.config.instance_id, + region_id=self.config.region_id, + manager_account=self.config.account, + manager_account_password=self.config.account_password, + namespace=self.config.namespace, + collection=self._collection_name, + dimension=embedding_dimension, + metrics=self.config.metrics, + metadata=metadata, + full_text_retrieval_fields=full_text_retrieval_fields, + ) + self._client.create_collection(request) + else: + raise ValueError(f"failed to create collection {self._collection_name}: {e}") + redis_client.set(collection_exist_cache_key, 1, ex=3600) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + from alibabacloud_gpdb20160503 import models as gpdb_20160503_models + + rows: list[gpdb_20160503_models.UpsertCollectionDataRequestRows] = [] + for doc, embedding in zip(documents, embeddings, strict=True): + if doc.metadata is not None: + metadata = { + "ref_doc_id": doc.metadata["doc_id"], + "page_content": doc.page_content, + "metadata_": json.dumps(doc.metadata), + } + rows.append( + gpdb_20160503_models.UpsertCollectionDataRequestRows( + vector=embedding, + metadata=metadata, + ) + ) + request = gpdb_20160503_models.UpsertCollectionDataRequest( + dbinstance_id=self.config.instance_id, + region_id=self.config.region_id, + namespace=self.config.namespace, + namespace_password=self.config.namespace_password, + collection=self._collection_name, + rows=rows, + ) + self._client.upsert_collection_data(request) + + def text_exists(self, id: str) -> bool: + from alibabacloud_gpdb20160503 import models as gpdb_20160503_models + + request = gpdb_20160503_models.QueryCollectionDataRequest( + dbinstance_id=self.config.instance_id, + region_id=self.config.region_id, + namespace=self.config.namespace, + namespace_password=self.config.namespace_password, + collection=self._collection_name, + metrics=self.config.metrics, + include_values=True, + vector=None, + content=None, + top_k=1, + filter=f"ref_doc_id='{id}'", + ) + response = self._client.query_collection_data(request) + return len(response.body.matches.match) > 0 + + def delete_by_ids(self, ids: list[str]) -> None: + from alibabacloud_gpdb20160503 import models as gpdb_20160503_models + + ids_str = ",".join(f"'{id}'" for id in ids) + ids_str = f"({ids_str})" + request = gpdb_20160503_models.DeleteCollectionDataRequest( + dbinstance_id=self.config.instance_id, + region_id=self.config.region_id, + namespace=self.config.namespace, + namespace_password=self.config.namespace_password, + collection=self._collection_name, + collection_data=None, + collection_data_filter=f"ref_doc_id IN {ids_str}", + ) + self._client.delete_collection_data(request) + + def delete_by_metadata_field(self, key: str, value: str) -> None: + from alibabacloud_gpdb20160503 import models as gpdb_20160503_models + + request = gpdb_20160503_models.DeleteCollectionDataRequest( + dbinstance_id=self.config.instance_id, + region_id=self.config.region_id, + namespace=self.config.namespace, + namespace_password=self.config.namespace_password, + collection=self._collection_name, + collection_data=None, + collection_data_filter=f"metadata_ ->> '{key}' = '{value}'", + ) + self._client.delete_collection_data(request) + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + from alibabacloud_gpdb20160503 import models as gpdb_20160503_models + + score_threshold = kwargs.get("score_threshold") or 0.0 + request = gpdb_20160503_models.QueryCollectionDataRequest( + dbinstance_id=self.config.instance_id, + region_id=self.config.region_id, + namespace=self.config.namespace, + namespace_password=self.config.namespace_password, + collection=self._collection_name, + include_values=kwargs.pop("include_values", True), + metrics=self.config.metrics, + vector=query_vector, + content=None, + top_k=kwargs.get("top_k", 4), + filter=None, + ) + response = self._client.query_collection_data(request) + documents = [] + for match in response.body.matches.match: + if match.score > score_threshold: + metadata = json.loads(match.metadata.get("metadata_")) + metadata["score"] = match.score + doc = Document( + page_content=match.metadata.get("page_content"), + vector=match.values.value, + metadata=metadata, + ) + documents.append(doc) + documents = sorted(documents, key=lambda x: x.metadata["score"] if x.metadata else 0, reverse=True) + return documents + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + from alibabacloud_gpdb20160503 import models as gpdb_20160503_models + + score_threshold = float(kwargs.get("score_threshold") or 0.0) + request = gpdb_20160503_models.QueryCollectionDataRequest( + dbinstance_id=self.config.instance_id, + region_id=self.config.region_id, + namespace=self.config.namespace, + namespace_password=self.config.namespace_password, + collection=self._collection_name, + include_values=kwargs.pop("include_values", True), + metrics=self.config.metrics, + vector=None, + content=query, + top_k=kwargs.get("top_k", 4), + filter=None, + ) + response = self._client.query_collection_data(request) + documents = [] + for match in response.body.matches.match: + if match.score > score_threshold: + metadata = json.loads(match.metadata.get("metadata_")) + metadata["score"] = match.score + doc = Document( + page_content=match.metadata.get("page_content"), + vector=match.values.value, + metadata=metadata, + ) + documents.append(doc) + documents = sorted(documents, key=lambda x: x.metadata["score"] if x.metadata else 0, reverse=True) + return documents + + def delete(self) -> None: + try: + from alibabacloud_gpdb20160503 import models as gpdb_20160503_models + + request = gpdb_20160503_models.DeleteCollectionRequest( + collection=self._collection_name, + dbinstance_id=self.config.instance_id, + namespace=self.config.namespace, + namespace_password=self.config.namespace_password, + region_id=self.config.region_id, + ) + self._client.delete_collection(request) + except Exception as e: + raise e diff --git a/api/core/rag/datasource/vdb/analyticdb/analyticdb_vector_sql.py b/api/core/rag/datasource/vdb/analyticdb/analyticdb_vector_sql.py new file mode 100644 index 0000000000000000000000000000000000000000..4d8f7929413cf249e79f9905a7374156c529ba91 --- /dev/null +++ b/api/core/rag/datasource/vdb/analyticdb/analyticdb_vector_sql.py @@ -0,0 +1,247 @@ +import json +import uuid +from contextlib import contextmanager +from typing import Any + +import psycopg2.extras # type: ignore +import psycopg2.pool # type: ignore +from pydantic import BaseModel, model_validator + +from core.rag.models.document import Document +from extensions.ext_redis import redis_client + + +class AnalyticdbVectorBySqlConfig(BaseModel): + host: str + port: int + account: str + account_password: str + min_connection: int + max_connection: int + namespace: str = "dify" + metrics: str = "cosine" + + @model_validator(mode="before") + @classmethod + def validate_config(cls, values: dict) -> dict: + if not values["host"]: + raise ValueError("config ANALYTICDB_HOST is required") + if not values["port"]: + raise ValueError("config ANALYTICDB_PORT is required") + if not values["account"]: + raise ValueError("config ANALYTICDB_ACCOUNT is required") + if not values["account_password"]: + raise ValueError("config ANALYTICDB_PASSWORD is required") + if not values["min_connection"]: + raise ValueError("config ANALYTICDB_MIN_CONNECTION is required") + if not values["max_connection"]: + raise ValueError("config ANALYTICDB_MAX_CONNECTION is required") + if values["min_connection"] > values["max_connection"]: + raise ValueError("config ANALYTICDB_MIN_CONNECTION should less than ANALYTICDB_MAX_CONNECTION") + return values + + +class AnalyticdbVectorBySql: + def __init__(self, collection_name: str, config: AnalyticdbVectorBySqlConfig): + self._collection_name = collection_name.lower() + self.databaseName = "knowledgebase" + self.config = config + self.table_name = f"{self.config.namespace}.{self._collection_name}" + self.pool = None + self._initialize() + if not self.pool: + self.pool = self._create_connection_pool() + + def _initialize(self) -> None: + cache_key = f"vector_initialize_{self.config.host}" + lock_name = f"{cache_key}_lock" + with redis_client.lock(lock_name, timeout=20): + database_exist_cache_key = f"vector_initialize_{self.config.host}" + if redis_client.get(database_exist_cache_key): + return + self._initialize_vector_database() + redis_client.set(database_exist_cache_key, 1, ex=3600) + + def _create_connection_pool(self): + return psycopg2.pool.SimpleConnectionPool( + self.config.min_connection, + self.config.max_connection, + host=self.config.host, + port=self.config.port, + user=self.config.account, + password=self.config.account_password, + database=self.databaseName, + ) + + @contextmanager + def _get_cursor(self): + assert self.pool is not None, "Connection pool is not initialized" + conn = self.pool.getconn() + cur = conn.cursor() + try: + yield cur + finally: + cur.close() + conn.commit() + self.pool.putconn(conn) + + def _initialize_vector_database(self) -> None: + conn = psycopg2.connect( + host=self.config.host, + port=self.config.port, + user=self.config.account, + password=self.config.account_password, + database="postgres", + ) + conn.autocommit = True + cur = conn.cursor() + try: + cur.execute(f"CREATE DATABASE {self.databaseName}") + except Exception as e: + if "already exists" in str(e): + return + raise e + finally: + cur.close() + conn.close() + self.pool = self._create_connection_pool() + with self._get_cursor() as cur: + try: + cur.execute("CREATE TEXT SEARCH CONFIGURATION zh_cn (PARSER = zhparser)") + cur.execute("ALTER TEXT SEARCH CONFIGURATION zh_cn ADD MAPPING FOR n,v,a,i,e,l,x WITH simple") + except Exception as e: + if "already exists" not in str(e): + raise e + cur.execute( + "CREATE OR REPLACE FUNCTION " + "public.to_tsquery_from_text(txt text, lang regconfig DEFAULT 'english'::regconfig) " + "RETURNS tsquery LANGUAGE sql IMMUTABLE STRICT AS $function$ " + "SELECT to_tsquery(lang, COALESCE(string_agg(split_part(word, ':', 1), ' | '), '')) " + "FROM (SELECT unnest(string_to_array(to_tsvector(lang, txt)::text, ' ')) AS word) " + "AS words_only;$function$" + ) + cur.execute(f"CREATE SCHEMA IF NOT EXISTS {self.config.namespace}") + + def _create_collection_if_not_exists(self, embedding_dimension: int): + cache_key = f"vector_indexing_{self._collection_name}" + lock_name = f"{cache_key}_lock" + with redis_client.lock(lock_name, timeout=20): + collection_exist_cache_key = f"vector_indexing_{self._collection_name}" + if redis_client.get(collection_exist_cache_key): + return + with self._get_cursor() as cur: + cur.execute( + f"CREATE TABLE IF NOT EXISTS {self.table_name}(" + f"id text PRIMARY KEY," + f"vector real[], ref_doc_id text, page_content text, metadata_ jsonb, " + f"to_tsvector TSVECTOR" + f") WITH (fillfactor=70) DISTRIBUTED BY (id);" + ) + if embedding_dimension is not None: + index_name = f"{self._collection_name}_embedding_idx" + cur.execute(f"ALTER TABLE {self.table_name} ALTER COLUMN vector SET STORAGE PLAIN") + cur.execute( + f"CREATE INDEX {index_name} ON {self.table_name} USING ann(vector) " + f"WITH(dim='{embedding_dimension}', distancemeasure='{self.config.metrics}', " + f"pq_enable=0, external_storage=0)" + ) + cur.execute(f"CREATE INDEX ON {self.table_name} USING gin(to_tsvector)") + redis_client.set(collection_exist_cache_key, 1, ex=3600) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + values = [] + id_prefix = str(uuid.uuid4()) + "_" + sql = f""" + INSERT INTO {self.table_name} + (id, ref_doc_id, vector, page_content, metadata_, to_tsvector) + VALUES (%s, %s, %s, %s, %s, to_tsvector('zh_cn', %s)); + """ + for i, doc in enumerate(documents): + if doc.metadata is not None: + values.append( + ( + id_prefix + str(i), + doc.metadata.get("doc_id", str(uuid.uuid4())), + embeddings[i], + doc.page_content, + json.dumps(doc.metadata), + doc.page_content, + ) + ) + with self._get_cursor() as cur: + psycopg2.extras.execute_batch(cur, sql, values) + + def text_exists(self, id: str) -> bool: + with self._get_cursor() as cur: + cur.execute(f"SELECT id FROM {self.table_name} WHERE ref_doc_id = %s", (id,)) + return cur.fetchone() is not None + + def delete_by_ids(self, ids: list[str]) -> None: + with self._get_cursor() as cur: + try: + cur.execute(f"DELETE FROM {self.table_name} WHERE ref_doc_id IN %s", (tuple(ids),)) + except Exception as e: + if "does not exist" not in str(e): + raise e + + def delete_by_metadata_field(self, key: str, value: str) -> None: + with self._get_cursor() as cur: + try: + cur.execute(f"DELETE FROM {self.table_name} WHERE metadata_->>%s = %s", (key, value)) + except Exception as e: + if "does not exist" not in str(e): + raise e + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + top_k = kwargs.get("top_k", 4) + score_threshold = float(kwargs.get("score_threshold") or 0.0) + with self._get_cursor() as cur: + query_vector_str = json.dumps(query_vector) + query_vector_str = "{" + query_vector_str[1:-1] + "}" + cur.execute( + f"SELECT t.id AS id, t.vector AS vector, (1.0 - t.score) AS score, " + f"t.page_content as page_content, t.metadata_ AS metadata_ " + f"FROM (SELECT id, vector, page_content, metadata_, vector <=> %s AS score " + f"FROM {self.table_name} ORDER BY score LIMIT {top_k} ) t", + (query_vector_str,), + ) + documents = [] + for record in cur: + id, vector, score, page_content, metadata = record + if score > score_threshold: + metadata["score"] = score + doc = Document( + page_content=page_content, + vector=vector, + metadata=metadata, + ) + documents.append(doc) + return documents + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + top_k = kwargs.get("top_k", 4) + with self._get_cursor() as cur: + cur.execute( + f"""SELECT id, vector, page_content, metadata_, + ts_rank(to_tsvector, to_tsquery_from_text(%s, 'zh_cn'), 32) AS score + FROM {self.table_name} + WHERE to_tsvector@@to_tsquery_from_text(%s, 'zh_cn') + ORDER BY score DESC + LIMIT {top_k}""", + (f"'{query}'", f"'{query}'"), + ) + documents = [] + for record in cur: + id, vector, page_content, metadata, score = record + metadata["score"] = score + doc = Document( + page_content=page_content, + vector=vector, + metadata=metadata, + ) + documents.append(doc) + return documents + + def delete(self) -> None: + with self._get_cursor() as cur: + cur.execute(f"DROP TABLE IF EXISTS {self.table_name}") diff --git a/api/core/rag/datasource/vdb/baidu/__init__.py b/api/core/rag/datasource/vdb/baidu/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/baidu/baidu_vector.py b/api/core/rag/datasource/vdb/baidu/baidu_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..a658495af75e4dd8b687accb7b6a54a3db0acd04 --- /dev/null +++ b/api/core/rag/datasource/vdb/baidu/baidu_vector.py @@ -0,0 +1,291 @@ +import json +import time +import uuid +from typing import Any + +import numpy as np +from pydantic import BaseModel, model_validator +from pymochow import MochowClient # type: ignore +from pymochow.auth.bce_credentials import BceCredentials # type: ignore +from pymochow.configuration import Configuration # type: ignore +from pymochow.exception import ServerError # type: ignore +from pymochow.model.enum import FieldType, IndexState, IndexType, MetricType, ServerErrCode, TableState # type: ignore +from pymochow.model.schema import Field, HNSWParams, Schema, VectorIndex # type: ignore +from pymochow.model.table import AnnSearch, HNSWSearchParams, Partition, Row # type: ignore + +from configs import dify_config +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from extensions.ext_redis import redis_client +from models.dataset import Dataset + + +class BaiduConfig(BaseModel): + endpoint: str + connection_timeout_in_mills: int = 30 * 1000 + account: str + api_key: str + database: str + index_type: str = "HNSW" + metric_type: str = "L2" + shard: int = 1 + replicas: int = 3 + + @model_validator(mode="before") + @classmethod + def validate_config(cls, values: dict) -> dict: + if not values["endpoint"]: + raise ValueError("config BAIDU_VECTOR_DB_ENDPOINT is required") + if not values["account"]: + raise ValueError("config BAIDU_VECTOR_DB_ACCOUNT is required") + if not values["api_key"]: + raise ValueError("config BAIDU_VECTOR_DB_API_KEY is required") + if not values["database"]: + raise ValueError("config BAIDU_VECTOR_DB_DATABASE is required") + return values + + +class BaiduVector(BaseVector): + field_id: str = "id" + field_vector: str = "vector" + field_text: str = "text" + field_metadata: str = "metadata" + field_app_id: str = "app_id" + field_annotation_id: str = "annotation_id" + index_vector: str = "vector_idx" + + def __init__(self, collection_name: str, config: BaiduConfig): + super().__init__(collection_name) + self._client_config = config + self._client = self._init_client(config) + self._db = self._init_database() + + def get_type(self) -> str: + return VectorType.BAIDU + + def to_index_struct(self) -> dict: + return {"type": self.get_type(), "vector_store": {"class_prefix": self._collection_name}} + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + self._create_table(len(embeddings[0])) + self.add_texts(texts, embeddings) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + texts = [doc.page_content for doc in documents] + metadatas = [doc.metadata for doc in documents if doc.metadata is not None] + total_count = len(documents) + batch_size = 1000 + + # upsert texts and embeddings batch by batch + table = self._db.table(self._collection_name) + for start in range(0, total_count, batch_size): + end = min(start + batch_size, total_count) + rows = [] + assert len(metadatas) == total_count, "metadatas length should be equal to total_count" + # FIXME do you need this assert? + for i in range(start, end, 1): + row = Row( + id=metadatas[i].get("doc_id", str(uuid.uuid4())), + vector=embeddings[i], + text=texts[i], + metadata=json.dumps(metadatas[i]), + app_id=metadatas[i].get("app_id", ""), + annotation_id=metadatas[i].get("annotation_id", ""), + ) + rows.append(row) + table.upsert(rows=rows) + + # rebuild vector index after upsert finished + table.rebuild_index(self.index_vector) + while True: + time.sleep(1) + index = table.describe_index(self.index_vector) + if index.state == IndexState.NORMAL: + break + + def text_exists(self, id: str) -> bool: + res = self._db.table(self._collection_name).query(primary_key={self.field_id: id}) + if res and res.code == 0: + return True + return False + + def delete_by_ids(self, ids: list[str]) -> None: + if not ids: + return + quoted_ids = [f"'{id}'" for id in ids] + self._db.table(self._collection_name).delete(filter=f"id IN({', '.join(quoted_ids)})") + + def delete_by_metadata_field(self, key: str, value: str) -> None: + self._db.table(self._collection_name).delete(filter=f"{key} = '{value}'") + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + query_vector = [float(val) if isinstance(val, np.float64) else val for val in query_vector] + anns = AnnSearch( + vector_field=self.field_vector, + vector_floats=query_vector, + params=HNSWSearchParams(ef=kwargs.get("ef", 10), limit=kwargs.get("top_k", 4)), + ) + res = self._db.table(self._collection_name).search( + anns=anns, + projections=[self.field_id, self.field_text, self.field_metadata], + retrieve_vector=True, + ) + score_threshold = float(kwargs.get("score_threshold") or 0.0) + return self._get_search_res(res, score_threshold) + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + # baidu vector database doesn't support bm25 search on current version + return [] + + def _get_search_res(self, res, score_threshold) -> list[Document]: + docs = [] + for row in res.rows: + row_data = row.get("row", {}) + meta = row_data.get(self.field_metadata) + if meta is not None: + meta = json.loads(meta) + score = row.get("score", 0.0) + if score > score_threshold: + meta["score"] = score + doc = Document(page_content=row_data.get(self.field_text), metadata=meta) + docs.append(doc) + + return docs + + def delete(self) -> None: + try: + self._db.drop_table(table_name=self._collection_name) + except ServerError as e: + if e.code == ServerErrCode.TABLE_NOT_EXIST: + pass + else: + raise + + def _init_client(self, config) -> MochowClient: + config = Configuration(credentials=BceCredentials(config.account, config.api_key), endpoint=config.endpoint) + client = MochowClient(config) + return client + + def _init_database(self): + exists = False + for db in self._client.list_databases(): + if db.database_name == self._client_config.database: + exists = True + break + # Create database if not existed + if exists: + return self._client.database(self._client_config.database) + else: + try: + self._client.create_database(database_name=self._client_config.database) + except ServerError as e: + if e.code == ServerErrCode.DB_ALREADY_EXIST: + pass + else: + raise + return + + def _table_existed(self) -> bool: + tables = self._db.list_table() + return any(table.table_name == self._collection_name for table in tables) + + def _create_table(self, dimension: int) -> None: + # Try to grab distributed lock and create table + lock_name = "vector_indexing_lock_{}".format(self._collection_name) + with redis_client.lock(lock_name, timeout=60): + table_exist_cache_key = "vector_indexing_{}".format(self._collection_name) + if redis_client.get(table_exist_cache_key): + return + + if self._table_existed(): + return + + self.delete() + + # check IndexType and MetricType + index_type = None + for k, v in IndexType.__members__.items(): + if k == self._client_config.index_type: + index_type = v + if index_type is None: + raise ValueError("unsupported index_type") + metric_type = None + for k, v in MetricType.__members__.items(): + if k == self._client_config.metric_type: + metric_type = v + if metric_type is None: + raise ValueError("unsupported metric_type") + + # Construct field schema + fields = [] + fields.append( + Field( + self.field_id, + FieldType.STRING, + primary_key=True, + partition_key=True, + auto_increment=False, + not_null=True, + ) + ) + fields.append(Field(self.field_metadata, FieldType.STRING, not_null=True)) + fields.append(Field(self.field_app_id, FieldType.STRING)) + fields.append(Field(self.field_annotation_id, FieldType.STRING)) + fields.append(Field(self.field_text, FieldType.TEXT, not_null=True)) + fields.append(Field(self.field_vector, FieldType.FLOAT_VECTOR, not_null=True, dimension=dimension)) + + # Construct vector index params + indexes = [] + indexes.append( + VectorIndex( + index_name="vector_idx", + index_type=index_type, + field="vector", + metric_type=metric_type, + params=HNSWParams(m=16, efconstruction=200), + ) + ) + + # Create table + self._db.create_table( + table_name=self._collection_name, + replication=self._client_config.replicas, + partition=Partition(partition_num=self._client_config.shard), + schema=Schema(fields=fields, indexes=indexes), + description="Table for Dify", + ) + + # Wait for table created + while True: + time.sleep(1) + table = self._db.describe_table(self._collection_name) + if table.state == TableState.NORMAL: + break + redis_client.set(table_exist_cache_key, 1, ex=3600) + + +class BaiduVectorFactory(AbstractVectorFactory): + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> BaiduVector: + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix.lower() + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id).lower() + dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.BAIDU, collection_name)) + + return BaiduVector( + collection_name=collection_name, + config=BaiduConfig( + endpoint=dify_config.BAIDU_VECTOR_DB_ENDPOINT or "", + connection_timeout_in_mills=dify_config.BAIDU_VECTOR_DB_CONNECTION_TIMEOUT_MS, + account=dify_config.BAIDU_VECTOR_DB_ACCOUNT or "", + api_key=dify_config.BAIDU_VECTOR_DB_API_KEY or "", + database=dify_config.BAIDU_VECTOR_DB_DATABASE or "", + shard=dify_config.BAIDU_VECTOR_DB_SHARD, + replicas=dify_config.BAIDU_VECTOR_DB_REPLICAS, + ), + ) diff --git a/api/core/rag/datasource/vdb/chroma/__init__.py b/api/core/rag/datasource/vdb/chroma/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/chroma/chroma_vector.py b/api/core/rag/datasource/vdb/chroma/chroma_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..907c4d22854de207c4ee9e1fa511299a40e2fa26 --- /dev/null +++ b/api/core/rag/datasource/vdb/chroma/chroma_vector.py @@ -0,0 +1,151 @@ +import json +from typing import Any, Optional + +import chromadb +from chromadb import QueryResult, Settings +from pydantic import BaseModel + +from configs import dify_config +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from extensions.ext_redis import redis_client +from models.dataset import Dataset + + +class ChromaConfig(BaseModel): + host: str + port: int + tenant: str + database: str + auth_provider: Optional[str] = None + auth_credentials: Optional[str] = None + + def to_chroma_params(self): + settings = Settings( + # auth + chroma_client_auth_provider=self.auth_provider, + chroma_client_auth_credentials=self.auth_credentials, + ) + + return { + "host": self.host, + "port": self.port, + "ssl": False, + "tenant": self.tenant, + "database": self.database, + "settings": settings, + } + + +class ChromaVector(BaseVector): + def __init__(self, collection_name: str, config: ChromaConfig): + super().__init__(collection_name) + self._client_config = config + self._client = chromadb.HttpClient(**self._client_config.to_chroma_params()) + + def get_type(self) -> str: + return VectorType.CHROMA + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + if texts: + # create collection + self.create_collection(self._collection_name) + + self.add_texts(texts, embeddings, **kwargs) + + def create_collection(self, collection_name: str): + lock_name = "vector_indexing_lock_{}".format(collection_name) + with redis_client.lock(lock_name, timeout=20): + collection_exist_cache_key = "vector_indexing_{}".format(self._collection_name) + if redis_client.get(collection_exist_cache_key): + return + self._client.get_or_create_collection(collection_name) + redis_client.set(collection_exist_cache_key, 1, ex=3600) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + uuids = self._get_uuids(documents) + texts = [d.page_content for d in documents] + metadatas = [d.metadata for d in documents] + + collection = self._client.get_or_create_collection(self._collection_name) + # FIXME: chromadb using numpy array, fix the type error later + collection.upsert(ids=uuids, documents=texts, embeddings=embeddings, metadatas=metadatas) # type: ignore + + def delete_by_metadata_field(self, key: str, value: str): + collection = self._client.get_or_create_collection(self._collection_name) + # FIXME: fix the type error later + collection.delete(where={key: {"$eq": value}}) # type: ignore + + def delete(self): + self._client.delete_collection(self._collection_name) + + def delete_by_ids(self, ids: list[str]) -> None: + if not ids: + return + collection = self._client.get_or_create_collection(self._collection_name) + collection.delete(ids=ids) + + def text_exists(self, id: str) -> bool: + collection = self._client.get_or_create_collection(self._collection_name) + response = collection.get(ids=[id]) + return len(response) > 0 + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + collection = self._client.get_or_create_collection(self._collection_name) + results: QueryResult = collection.query(query_embeddings=query_vector, n_results=kwargs.get("top_k", 4)) + score_threshold = float(kwargs.get("score_threshold") or 0.0) + + # Check if results contain data + if not results["ids"] or not results["documents"] or not results["metadatas"] or not results["distances"]: + return [] + + ids = results["ids"][0] + documents = results["documents"][0] + metadatas = results["metadatas"][0] + distances = results["distances"][0] + + docs = [] + for index in range(len(ids)): + distance = distances[index] + metadata = dict(metadatas[index]) + if distance >= score_threshold: + metadata["score"] = distance + doc = Document( + page_content=documents[index], + metadata=metadata, + ) + docs.append(doc) + # Sort the documents by score in descending order + docs = sorted(docs, key=lambda x: x.metadata["score"] if x.metadata is not None else 0, reverse=True) + return docs + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + # chroma does not support BM25 full text searching + return [] + + +class ChromaVectorFactory(AbstractVectorFactory): + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> BaseVector: + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix.lower() + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id).lower() + index_struct_dict = {"type": VectorType.CHROMA, "vector_store": {"class_prefix": collection_name}} + dataset.index_struct = json.dumps(index_struct_dict) + + return ChromaVector( + collection_name=collection_name, + config=ChromaConfig( + host=dify_config.CHROMA_HOST or "", + port=dify_config.CHROMA_PORT, + tenant=dify_config.CHROMA_TENANT or chromadb.DEFAULT_TENANT, + database=dify_config.CHROMA_DATABASE or chromadb.DEFAULT_DATABASE, + auth_provider=dify_config.CHROMA_AUTH_PROVIDER, + auth_credentials=dify_config.CHROMA_AUTH_CREDENTIALS, + ), + ) diff --git a/api/core/rag/datasource/vdb/couchbase/__init__.py b/api/core/rag/datasource/vdb/couchbase/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/couchbase/couchbase_vector.py b/api/core/rag/datasource/vdb/couchbase/couchbase_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..68a9952789e5b6e4faa0e374dfd5452710f4ab15 --- /dev/null +++ b/api/core/rag/datasource/vdb/couchbase/couchbase_vector.py @@ -0,0 +1,378 @@ +import json +import logging +import time +import uuid +from datetime import timedelta +from typing import Any + +from couchbase import search # type: ignore +from couchbase.auth import PasswordAuthenticator # type: ignore +from couchbase.cluster import Cluster # type: ignore +from couchbase.management.search import SearchIndex # type: ignore + +# needed for options -- cluster, timeout, SQL++ (N1QL) query, etc. +from couchbase.options import ClusterOptions, SearchOptions # type: ignore +from couchbase.vector_search import VectorQuery, VectorSearch # type: ignore +from flask import current_app +from pydantic import BaseModel, model_validator + +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from extensions.ext_redis import redis_client +from models.dataset import Dataset + +logger = logging.getLogger(__name__) + + +class CouchbaseConfig(BaseModel): + connection_string: str + user: str + password: str + bucket_name: str + scope_name: str + + @model_validator(mode="before") + @classmethod + def validate_config(cls, values: dict) -> dict: + if not values.get("connection_string"): + raise ValueError("config COUCHBASE_CONNECTION_STRING is required") + if not values.get("user"): + raise ValueError("config COUCHBASE_USER is required") + if not values.get("password"): + raise ValueError("config COUCHBASE_PASSWORD is required") + if not values.get("bucket_name"): + raise ValueError("config COUCHBASE_PASSWORD is required") + if not values.get("scope_name"): + raise ValueError("config COUCHBASE_SCOPE_NAME is required") + return values + + +class CouchbaseVector(BaseVector): + def __init__(self, collection_name: str, config: CouchbaseConfig): + super().__init__(collection_name) + self._client_config = config + + """Connect to couchbase""" + + auth = PasswordAuthenticator(config.user, config.password) + options = ClusterOptions(auth) + self._cluster = Cluster(config.connection_string, options) + self._bucket = self._cluster.bucket(config.bucket_name) + self._scope = self._bucket.scope(config.scope_name) + self._bucket_name = config.bucket_name + self._scope_name = config.scope_name + + # Wait until the cluster is ready for use. + self._cluster.wait_until_ready(timedelta(seconds=5)) + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + index_id = str(uuid.uuid4()).replace("-", "") + self._create_collection(uuid=index_id, vector_length=len(embeddings[0])) + self.add_texts(texts, embeddings) + + def _create_collection(self, vector_length: int, uuid: str): + lock_name = "vector_indexing_lock_{}".format(self._collection_name) + with redis_client.lock(lock_name, timeout=20): + collection_exist_cache_key = "vector_indexing_{}".format(self._collection_name) + if redis_client.get(collection_exist_cache_key): + return + if self._collection_exists(self._collection_name): + return + manager = self._bucket.collections() + manager.create_collection(self._client_config.scope_name, self._collection_name) + + index_manager = self._scope.search_indexes() + + index_definition = json.loads(""" +{ + "type": "fulltext-index", + "name": "Embeddings._default.Vector_Search", + "uuid": "26d4db528e78b716", + "sourceType": "gocbcore", + "sourceName": "Embeddings", + "sourceUUID": "2242e4a25b4decd6650c9c7b3afa1dbf", + "planParams": { + "maxPartitionsPerPIndex": 1024, + "indexPartitions": 1 + }, + "params": { + "doc_config": { + "docid_prefix_delim": "", + "docid_regexp": "", + "mode": "scope.collection.type_field", + "type_field": "type" + }, + "mapping": { + "analysis": { }, + "default_analyzer": "standard", + "default_datetime_parser": "dateTimeOptional", + "default_field": "_all", + "default_mapping": { + "dynamic": true, + "enabled": true + }, + "default_type": "_default", + "docvalues_dynamic": false, + "index_dynamic": true, + "store_dynamic": true, + "type_field": "_type", + "types": { + "collection_name": { + "dynamic": true, + "enabled": true, + "properties": { + "embedding": { + "dynamic": false, + "enabled": true, + "fields": [ + { + "dims": 1536, + "index": true, + "name": "embedding", + "similarity": "dot_product", + "type": "vector", + "vector_index_optimized_for": "recall" + } + ] + }, + "metadata": { + "dynamic": true, + "enabled": true + }, + "text": { + "dynamic": false, + "enabled": true, + "fields": [ + { + "index": true, + "name": "text", + "store": true, + "type": "text" + } + ] + } + } + } + } + }, + "store": { + "indexType": "scorch", + "segmentVersion": 16 + } + }, + "sourceParams": { } + } +""") + index_definition["name"] = self._collection_name + "_search" + index_definition["uuid"] = uuid + index_definition["params"]["mapping"]["types"]["collection_name"]["properties"]["embedding"]["fields"][0][ + "dims" + ] = vector_length + index_definition["params"]["mapping"]["types"][self._scope_name + "." + self._collection_name] = ( + index_definition["params"]["mapping"]["types"].pop("collection_name") + ) + time.sleep(2) + index_manager.upsert_index( + SearchIndex( + index_definition["name"], + params=index_definition["params"], + source_name=self._bucket_name, + ), + ) + time.sleep(1) + + redis_client.set(collection_exist_cache_key, 1, ex=3600) + + def _collection_exists(self, name: str): + scope_collection_map: dict[str, Any] = {} + + # Get a list of all scopes in the bucket + for scope in self._bucket.collections().get_all_scopes(): + scope_collection_map[scope.name] = [] + + # Get a list of all the collections in the scope + for collection in scope.collections: + scope_collection_map[scope.name].append(collection.name) + + # Check if the collection exists in the scope + return self._collection_name in scope_collection_map[self._scope_name] + + def get_type(self) -> str: + return VectorType.COUCHBASE + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + uuids = self._get_uuids(documents) + texts = [d.page_content for d in documents] + metadatas = [d.metadata for d in documents] + + doc_ids = [] + + documents_to_insert = [ + {"text": text, "embedding": vector, "metadata": metadata} + for id, text, vector, metadata in zip(uuids, texts, embeddings, metadatas) + ] + for doc, id in zip(documents_to_insert, uuids): + result = self._scope.collection(self._collection_name).upsert(id, doc) + + doc_ids.extend(uuids) + + return doc_ids + + def text_exists(self, id: str) -> bool: + # Use a parameterized query for safety and correctness + query = f""" + SELECT COUNT(1) AS count FROM + `{self._client_config.bucket_name}`.{self._client_config.scope_name}.{self._collection_name} + WHERE META().id = $doc_id + """ + # Pass the id as a parameter to the query + result = self._cluster.query(query, named_parameters={"doc_id": id}).execute() + for row in result: + return bool(row["count"] > 0) + return False # Return False if no rows are returned + + def delete_by_ids(self, ids: list[str]) -> None: + query = f""" + DELETE FROM `{self._bucket_name}`.{self._client_config.scope_name}.{self._collection_name} + WHERE META().id IN $doc_ids; + """ + try: + self._cluster.query(query, named_parameters={"doc_ids": ids}).execute() + except Exception as e: + logger.exception(f"Failed to delete documents, ids: {ids}") + + def delete_by_document_id(self, document_id: str): + query = f""" + DELETE FROM + `{self._client_config.bucket_name}`.{self._client_config.scope_name}.{self._collection_name} + WHERE META().id = $doc_id; + """ + self._cluster.query(query, named_parameters={"doc_id": document_id}).execute() + + # def get_ids_by_metadata_field(self, key: str, value: str): + # query = f""" + # SELECT id FROM + # `{self._client_config.bucket_name}`.{self._client_config.scope_name}.{self._collection_name} + # WHERE `metadata.{key}` = $value; + # """ + # result = self._cluster.query(query, named_parameters={'value':value}) + # return [row['id'] for row in result.rows()] + + def delete_by_metadata_field(self, key: str, value: str) -> None: + query = f""" + DELETE FROM `{self._client_config.bucket_name}`.{self._client_config.scope_name}.{self._collection_name} + WHERE metadata.{key} = $value; + """ + self._cluster.query(query, named_parameters={"value": value}).execute() + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + top_k = kwargs.get("top_k", 5) + score_threshold = kwargs.get("score_threshold") or 0.0 + + search_req = search.SearchRequest.create( + VectorSearch.from_vector_query( + VectorQuery( + "embedding", + query_vector, + top_k, + ) + ) + ) + try: + search_iter = self._scope.search( + self._collection_name + "_search", + search_req, + SearchOptions(limit=top_k, collections=[self._collection_name], fields=["*"]), + ) + + docs = [] + # Parse the results + for row in search_iter.rows(): + text = row.fields.pop("text") + metadata = self._format_metadata(row.fields) + score = row.score + metadata["score"] = score + doc = Document(page_content=text, metadata=metadata) + if score >= score_threshold: + docs.append(doc) + except Exception as e: + raise ValueError(f"Search failed with error: {e}") + + return docs + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + top_k = kwargs.get("top_k", 2) + try: + CBrequest = search.SearchRequest.create(search.QueryStringQuery("text:" + query)) + search_iter = self._scope.search( + self._collection_name + "_search", CBrequest, SearchOptions(limit=top_k, fields=["*"]) + ) + + docs = [] + for row in search_iter.rows(): + text = row.fields.pop("text") + metadata = self._format_metadata(row.fields) + score = row.score + metadata["score"] = score + doc = Document(page_content=text, metadata=metadata) + docs.append(doc) + + except Exception as e: + raise ValueError(f"Search failed with error: {e}") + + return docs + + def delete(self): + manager = self._bucket.collections() + scopes = manager.get_all_scopes() + + for scope in scopes: + for collection in scope.collections: + if collection.name == self._collection_name: + manager.drop_collection("_default", self._collection_name) + + def _format_metadata(self, row_fields: dict[str, Any]) -> dict[str, Any]: + """Helper method to format the metadata from the Couchbase Search API. + Args: + row_fields (Dict[str, Any]): The fields to format. + + Returns: + Dict[str, Any]: The formatted metadata. + """ + metadata = {} + for key, value in row_fields.items(): + # Couchbase Search returns the metadata key with a prefix + # `metadata.` We remove it to get the original metadata key + if key.startswith("metadata"): + new_key = key.split("metadata" + ".")[-1] + metadata[new_key] = value + else: + metadata[key] = value + + return metadata + + +class CouchbaseVectorFactory(AbstractVectorFactory): + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> CouchbaseVector: + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id) + dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.COUCHBASE, collection_name)) + + config = current_app.config + return CouchbaseVector( + collection_name=collection_name, + config=CouchbaseConfig( + connection_string=config.get("COUCHBASE_CONNECTION_STRING", ""), + user=config.get("COUCHBASE_USER", ""), + password=config.get("COUCHBASE_PASSWORD", ""), + bucket_name=config.get("COUCHBASE_BUCKET_NAME", ""), + scope_name=config.get("COUCHBASE_SCOPE_NAME", ""), + ), + ) diff --git a/api/core/rag/datasource/vdb/elasticsearch/__init__.py b/api/core/rag/datasource/vdb/elasticsearch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/elasticsearch/elasticsearch_ja_vector.py b/api/core/rag/datasource/vdb/elasticsearch/elasticsearch_ja_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..27575197faccfe84bce034980a653ef5052a5437 --- /dev/null +++ b/api/core/rag/datasource/vdb/elasticsearch/elasticsearch_ja_vector.py @@ -0,0 +1,104 @@ +import json +import logging +from typing import Any, Optional + +from flask import current_app + +from core.rag.datasource.vdb.elasticsearch.elasticsearch_vector import ( + ElasticSearchConfig, + ElasticSearchVector, + ElasticSearchVectorFactory, +) +from core.rag.datasource.vdb.field import Field +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from extensions.ext_redis import redis_client +from models.dataset import Dataset + +logger = logging.getLogger(__name__) + + +class ElasticSearchJaVector(ElasticSearchVector): + def create_collection( + self, + embeddings: list[list[float]], + metadatas: Optional[list[dict[Any, Any]]] = None, + index_params: Optional[dict] = None, + ): + lock_name = f"vector_indexing_lock_{self._collection_name}" + with redis_client.lock(lock_name, timeout=20): + collection_exist_cache_key = f"vector_indexing_{self._collection_name}" + if redis_client.get(collection_exist_cache_key): + logger.info(f"Collection {self._collection_name} already exists.") + return + + if not self._client.indices.exists(index=self._collection_name): + dim = len(embeddings[0]) + settings = { + "analysis": { + "analyzer": { + "ja_analyzer": { + "type": "custom", + "char_filter": [ + "icu_normalizer", + "kuromoji_iteration_mark", + ], + "tokenizer": "kuromoji_tokenizer", + "filter": [ + "kuromoji_baseform", + "kuromoji_part_of_speech", + "ja_stop", + "kuromoji_number", + "kuromoji_stemmer", + ], + } + } + } + } + mappings = { + "properties": { + Field.CONTENT_KEY.value: { + "type": "text", + "analyzer": "ja_analyzer", + "search_analyzer": "ja_analyzer", + }, + Field.VECTOR.value: { # Make sure the dimension is correct here + "type": "dense_vector", + "dims": dim, + "index": True, + "similarity": "cosine", + }, + Field.METADATA_KEY.value: { + "type": "object", + "properties": { + "doc_id": {"type": "keyword"} # Map doc_id to keyword type + }, + }, + } + } + self._client.indices.create(index=self._collection_name, settings=settings, mappings=mappings) + + redis_client.set(collection_exist_cache_key, 1, ex=3600) + + +class ElasticSearchJaVectorFactory(ElasticSearchVectorFactory): + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> ElasticSearchJaVector: + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id) + dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.ELASTICSEARCH, collection_name)) + + config = current_app.config + return ElasticSearchJaVector( + index_name=collection_name, + config=ElasticSearchConfig( + host=config.get("ELASTICSEARCH_HOST", "localhost"), + port=config.get("ELASTICSEARCH_PORT", 9200), + username=config.get("ELASTICSEARCH_USERNAME", ""), + password=config.get("ELASTICSEARCH_PASSWORD", ""), + ), + attributes=[], + ) diff --git a/api/core/rag/datasource/vdb/elasticsearch/elasticsearch_vector.py b/api/core/rag/datasource/vdb/elasticsearch/elasticsearch_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..cca696baeec94a608ea66ec27511986d8fa16d72 --- /dev/null +++ b/api/core/rag/datasource/vdb/elasticsearch/elasticsearch_vector.py @@ -0,0 +1,223 @@ +import json +import logging +import math +from typing import Any, Optional, cast +from urllib.parse import urlparse + +import requests +from elasticsearch import Elasticsearch +from flask import current_app +from pydantic import BaseModel, model_validator + +from core.rag.datasource.vdb.field import Field +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from extensions.ext_redis import redis_client +from models.dataset import Dataset + +logger = logging.getLogger(__name__) + + +class ElasticSearchConfig(BaseModel): + host: str + port: int + username: str + password: str + + @model_validator(mode="before") + @classmethod + def validate_config(cls, values: dict) -> dict: + if not values["host"]: + raise ValueError("config HOST is required") + if not values["port"]: + raise ValueError("config PORT is required") + if not values["username"]: + raise ValueError("config USERNAME is required") + if not values["password"]: + raise ValueError("config PASSWORD is required") + return values + + +class ElasticSearchVector(BaseVector): + def __init__(self, index_name: str, config: ElasticSearchConfig, attributes: list): + super().__init__(index_name.lower()) + self._client = self._init_client(config) + self._version = self._get_version() + self._check_version() + self._attributes = attributes + + def _init_client(self, config: ElasticSearchConfig) -> Elasticsearch: + try: + parsed_url = urlparse(config.host) + if parsed_url.scheme in {"http", "https"}: + hosts = f"{config.host}:{config.port}" + else: + hosts = f"http://{config.host}:{config.port}" + client = Elasticsearch( + hosts=hosts, + basic_auth=(config.username, config.password), + request_timeout=100000, + retry_on_timeout=True, + max_retries=10000, + ) + except requests.exceptions.ConnectionError: + raise ConnectionError("Vector database connection error") + + return client + + def _get_version(self) -> str: + info = self._client.info() + return cast(str, info["version"]["number"]) + + def _check_version(self): + if self._version < "8.0.0": + raise ValueError("Elasticsearch vector database version must be greater than 8.0.0") + + def get_type(self) -> str: + return VectorType.ELASTICSEARCH + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + uuids = self._get_uuids(documents) + for i in range(len(documents)): + self._client.index( + index=self._collection_name, + id=uuids[i], + document={ + Field.CONTENT_KEY.value: documents[i].page_content, + Field.VECTOR.value: embeddings[i] or None, + Field.METADATA_KEY.value: documents[i].metadata or {}, + }, + ) + self._client.indices.refresh(index=self._collection_name) + return uuids + + def text_exists(self, id: str) -> bool: + return bool(self._client.exists(index=self._collection_name, id=id)) + + def delete_by_ids(self, ids: list[str]) -> None: + if not ids: + return + for id in ids: + self._client.delete(index=self._collection_name, id=id) + + def delete_by_metadata_field(self, key: str, value: str) -> None: + query_str = {"query": {"match": {f"metadata.{key}": f"{value}"}}} + results = self._client.search(index=self._collection_name, body=query_str) + ids = [hit["_id"] for hit in results["hits"]["hits"]] + if ids: + self.delete_by_ids(ids) + + def delete(self) -> None: + self._client.indices.delete(index=self._collection_name) + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + top_k = kwargs.get("top_k", 4) + num_candidates = math.ceil(top_k * 1.5) + knn = {"field": Field.VECTOR.value, "query_vector": query_vector, "k": top_k, "num_candidates": num_candidates} + + results = self._client.search(index=self._collection_name, knn=knn, size=top_k) + + docs_and_scores = [] + for hit in results["hits"]["hits"]: + docs_and_scores.append( + ( + Document( + page_content=hit["_source"][Field.CONTENT_KEY.value], + vector=hit["_source"][Field.VECTOR.value], + metadata=hit["_source"][Field.METADATA_KEY.value], + ), + hit["_score"], + ) + ) + + docs = [] + for doc, score in docs_and_scores: + score_threshold = float(kwargs.get("score_threshold") or 0.0) + if score > score_threshold: + if doc.metadata is not None: + doc.metadata["score"] = score + docs.append(doc) + + return docs + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + query_str = {"match": {Field.CONTENT_KEY.value: query}} + results = self._client.search(index=self._collection_name, query=query_str, size=kwargs.get("top_k", 4)) + docs = [] + for hit in results["hits"]["hits"]: + docs.append( + Document( + page_content=hit["_source"][Field.CONTENT_KEY.value], + vector=hit["_source"][Field.VECTOR.value], + metadata=hit["_source"][Field.METADATA_KEY.value], + ) + ) + + return docs + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + metadatas = [d.metadata if d.metadata is not None else {} for d in texts] + self.create_collection(embeddings, metadatas) + self.add_texts(texts, embeddings, **kwargs) + + def create_collection( + self, + embeddings: list[list[float]], + metadatas: Optional[list[dict[Any, Any]]] = None, + index_params: Optional[dict] = None, + ): + lock_name = f"vector_indexing_lock_{self._collection_name}" + with redis_client.lock(lock_name, timeout=20): + collection_exist_cache_key = f"vector_indexing_{self._collection_name}" + if redis_client.get(collection_exist_cache_key): + logger.info(f"Collection {self._collection_name} already exists.") + return + + if not self._client.indices.exists(index=self._collection_name): + dim = len(embeddings[0]) + mappings = { + "properties": { + Field.CONTENT_KEY.value: {"type": "text"}, + Field.VECTOR.value: { # Make sure the dimension is correct here + "type": "dense_vector", + "dims": dim, + "index": True, + "similarity": "cosine", + }, + Field.METADATA_KEY.value: { + "type": "object", + "properties": { + "doc_id": {"type": "keyword"} # Map doc_id to keyword type + }, + }, + } + } + self._client.indices.create(index=self._collection_name, mappings=mappings) + + redis_client.set(collection_exist_cache_key, 1, ex=3600) + + +class ElasticSearchVectorFactory(AbstractVectorFactory): + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> ElasticSearchVector: + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id) + dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.ELASTICSEARCH, collection_name)) + + config = current_app.config + return ElasticSearchVector( + index_name=collection_name, + config=ElasticSearchConfig( + host=config.get("ELASTICSEARCH_HOST", "localhost"), + port=config.get("ELASTICSEARCH_PORT", 9200), + username=config.get("ELASTICSEARCH_USERNAME", ""), + password=config.get("ELASTICSEARCH_PASSWORD", ""), + ), + attributes=[], + ) diff --git a/api/core/rag/datasource/vdb/field.py b/api/core/rag/datasource/vdb/field.py new file mode 100644 index 0000000000000000000000000000000000000000..a64407bce1c14a77573d9bc5b0a3a65437944954 --- /dev/null +++ b/api/core/rag/datasource/vdb/field.py @@ -0,0 +1,13 @@ +from enum import Enum + + +class Field(Enum): + CONTENT_KEY = "page_content" + METADATA_KEY = "metadata" + GROUP_KEY = "group_id" + VECTOR = "vector" + # Sparse Vector aims to support full text search + SPARSE_VECTOR = "sparse_vector" + TEXT_KEY = "text" + PRIMARY_KEY = "id" + DOC_ID = "metadata.doc_id" diff --git a/api/core/rag/datasource/vdb/lindorm/__init__.py b/api/core/rag/datasource/vdb/lindorm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/lindorm/lindorm_vector.py b/api/core/rag/datasource/vdb/lindorm/lindorm_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..66fba763e731c199acec3a901ee5f7b30a3eeeb4 --- /dev/null +++ b/api/core/rag/datasource/vdb/lindorm/lindorm_vector.py @@ -0,0 +1,509 @@ +import copy +import json +import logging +from typing import Any, Optional + +from opensearchpy import OpenSearch +from pydantic import BaseModel, model_validator + +from configs import dify_config +from core.rag.datasource.vdb.field import Field +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from extensions.ext_redis import redis_client +from models.dataset import Dataset + +logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") +logging.getLogger("lindorm").setLevel(logging.WARN) + +ROUTING_FIELD = "routing_field" +UGC_INDEX_PREFIX = "ugc_index" + + +class LindormVectorStoreConfig(BaseModel): + hosts: str + username: Optional[str] = None + password: Optional[str] = None + using_ugc: Optional[bool] = False + + @model_validator(mode="before") + @classmethod + def validate_config(cls, values: dict) -> dict: + if not values["hosts"]: + raise ValueError("config URL is required") + if not values["username"]: + raise ValueError("config USERNAME is required") + if not values["password"]: + raise ValueError("config PASSWORD is required") + return values + + def to_opensearch_params(self) -> dict[str, Any]: + params: dict[str, Any] = {"hosts": self.hosts} + if self.username and self.password: + params["http_auth"] = (self.username, self.password) + return params + + +class LindormVectorStore(BaseVector): + def __init__(self, collection_name: str, config: LindormVectorStoreConfig, using_ugc: bool, **kwargs): + self._routing = None + self._routing_field = None + if using_ugc: + routing_value: str | None = kwargs.get("routing_value") + if routing_value is None: + raise ValueError("UGC index should init vector with valid 'routing_value' parameter value") + self._routing = routing_value.lower() + self._routing_field = ROUTING_FIELD + ugc_index_name = collection_name + super().__init__(ugc_index_name.lower()) + else: + super().__init__(collection_name.lower()) + self._client_config = config + self._client = OpenSearch(**config.to_opensearch_params()) + self._using_ugc = using_ugc + self.kwargs = kwargs + + def get_type(self) -> str: + return VectorType.LINDORM + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + self.create_collection(len(embeddings[0]), **kwargs) + self.add_texts(texts, embeddings) + + def refresh(self): + self._client.indices.refresh(index=self._collection_name) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + actions = [] + uuids = self._get_uuids(documents) + for i in range(len(documents)): + action_header = { + "index": { + "_index": self.collection_name.lower(), + "_id": uuids[i], + } + } + action_values: dict[str, Any] = { + Field.CONTENT_KEY.value: documents[i].page_content, + Field.VECTOR.value: embeddings[i], # Make sure you pass an array here + Field.METADATA_KEY.value: documents[i].metadata, + } + if self._using_ugc: + action_header["index"]["routing"] = self._routing + if self._routing_field is not None: + action_values[self._routing_field] = self._routing + actions.append(action_header) + actions.append(action_values) + response = self._client.bulk(actions) + if response["errors"]: + for item in response["items"]: + print(f"{item['index']['status']}: {item['index']['error']['type']}") + else: + self.refresh() + + def get_ids_by_metadata_field(self, key: str, value: str): + query: dict[str, Any] = { + "query": {"bool": {"must": [{"term": {f"{Field.METADATA_KEY.value}.{key}.keyword": value}}]}} + } + if self._using_ugc: + query["query"]["bool"]["must"].append({"term": {f"{self._routing_field}.keyword": self._routing}}) + response = self._client.search(index=self._collection_name, body=query) + if response["hits"]["hits"]: + return [hit["_id"] for hit in response["hits"]["hits"]] + else: + return None + + def delete_by_metadata_field(self, key: str, value: str): + ids = self.get_ids_by_metadata_field(key, value) + if ids: + self.delete_by_ids(ids) + + def delete_by_ids(self, ids: list[str]) -> None: + params = {} + if self._using_ugc: + params["routing"] = self._routing + for id in ids: + if self._client.exists(index=self._collection_name, id=id, params=params): + params = {} + if self._using_ugc: + params["routing"] = self._routing + self._client.delete(index=self._collection_name, id=id, params=params) + self.refresh() + else: + logger.warning(f"DELETE BY ID: ID {id} does not exist in the index.") + + def delete(self) -> None: + if self._using_ugc: + routing_filter_query = { + "query": {"bool": {"must": [{"term": {f"{self._routing_field}.keyword": self._routing}}]}} + } + self._client.delete_by_query(self._collection_name, body=routing_filter_query) + self.refresh() + else: + if self._client.indices.exists(index=self._collection_name): + self._client.indices.delete(index=self._collection_name, params={"timeout": 60}) + logger.info("Delete index success") + else: + logger.warning(f"Index '{self._collection_name}' does not exist. No deletion performed.") + + def text_exists(self, id: str) -> bool: + try: + params = {} + if self._using_ugc: + params["routing"] = self._routing + self._client.get(index=self._collection_name, id=id, params=params) + return True + except: + return False + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + if not isinstance(query_vector, list): + raise ValueError("query_vector should be a list of floats") + + if not all(isinstance(x, float) for x in query_vector): + raise ValueError("All elements in query_vector should be floats") + + top_k = kwargs.get("top_k", 10) + query = default_vector_search_query(query_vector=query_vector, k=top_k, **kwargs) + try: + params = {} + if self._using_ugc: + params["routing"] = self._routing + response = self._client.search(index=self._collection_name, body=query, params=params) + except Exception as e: + logger.exception(f"Error executing vector search, query: {query}") + raise + + docs_and_scores = [] + for hit in response["hits"]["hits"]: + docs_and_scores.append( + ( + Document( + page_content=hit["_source"][Field.CONTENT_KEY.value], + vector=hit["_source"][Field.VECTOR.value], + metadata=hit["_source"][Field.METADATA_KEY.value], + ), + hit["_score"], + ) + ) + docs = [] + for doc, score in docs_and_scores: + score_threshold = kwargs.get("score_threshold", 0.0) or 0.0 + if score > score_threshold: + if doc.metadata is not None: + doc.metadata["score"] = score + docs.append(doc) + + return docs + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + must = kwargs.get("must") + must_not = kwargs.get("must_not") + should = kwargs.get("should") + minimum_should_match = kwargs.get("minimum_should_match", 0) + top_k = kwargs.get("top_k", 10) + filters = kwargs.get("filter") + routing = self._routing + full_text_query = default_text_search_query( + query_text=query, + k=top_k, + text_field=Field.CONTENT_KEY.value, + must=must, + must_not=must_not, + should=should, + minimum_should_match=minimum_should_match, + filters=filters, + routing=routing, + routing_field=self._routing_field, + ) + response = self._client.search(index=self._collection_name, body=full_text_query) + docs = [] + for hit in response["hits"]["hits"]: + docs.append( + Document( + page_content=hit["_source"][Field.CONTENT_KEY.value], + vector=hit["_source"][Field.VECTOR.value], + metadata=hit["_source"][Field.METADATA_KEY.value], + ) + ) + + return docs + + def create_collection(self, dimension: int, **kwargs): + lock_name = f"vector_indexing_lock_{self._collection_name}" + with redis_client.lock(lock_name, timeout=20): + collection_exist_cache_key = f"vector_indexing_{self._collection_name}" + if redis_client.get(collection_exist_cache_key): + logger.info(f"Collection {self._collection_name} already exists.") + return + if self._client.indices.exists(index=self._collection_name): + logger.info(f"{self._collection_name.lower()} already exists.") + redis_client.set(collection_exist_cache_key, 1, ex=3600) + return + if len(self.kwargs) == 0 and len(kwargs) != 0: + self.kwargs = copy.deepcopy(kwargs) + vector_field = kwargs.pop("vector_field", Field.VECTOR.value) + shards = kwargs.pop("shards", 4) + + engine = kwargs.pop("engine", "lvector") + method_name = kwargs.pop("method_name", dify_config.DEFAULT_INDEX_TYPE) + space_type = kwargs.pop("space_type", dify_config.DEFAULT_DISTANCE_TYPE) + data_type = kwargs.pop("data_type", "float") + + hnsw_m = kwargs.pop("hnsw_m", 24) + hnsw_ef_construction = kwargs.pop("hnsw_ef_construction", 500) + ivfpq_m = kwargs.pop("ivfpq_m", dimension) + nlist = kwargs.pop("nlist", 1000) + centroids_use_hnsw = kwargs.pop("centroids_use_hnsw", nlist >= 5000) + centroids_hnsw_m = kwargs.pop("centroids_hnsw_m", 24) + centroids_hnsw_ef_construct = kwargs.pop("centroids_hnsw_ef_construct", 500) + centroids_hnsw_ef_search = kwargs.pop("centroids_hnsw_ef_search", 100) + mapping = default_text_mapping( + dimension, + method_name, + space_type=space_type, + shards=shards, + engine=engine, + data_type=data_type, + vector_field=vector_field, + hnsw_m=hnsw_m, + hnsw_ef_construction=hnsw_ef_construction, + nlist=nlist, + ivfpq_m=ivfpq_m, + centroids_use_hnsw=centroids_use_hnsw, + centroids_hnsw_m=centroids_hnsw_m, + centroids_hnsw_ef_construct=centroids_hnsw_ef_construct, + centroids_hnsw_ef_search=centroids_hnsw_ef_search, + using_ugc=self._using_ugc, + **kwargs, + ) + self._client.indices.create(index=self._collection_name.lower(), body=mapping) + redis_client.set(collection_exist_cache_key, 1, ex=3600) + # logger.info(f"create index success: {self._collection_name}") + + +def default_text_mapping(dimension: int, method_name: str, **kwargs: Any) -> dict: + excludes_from_source = kwargs.get("excludes_from_source") + analyzer = kwargs.get("analyzer", "ik_max_word") + text_field = kwargs.get("text_field", Field.CONTENT_KEY.value) + engine = kwargs["engine"] + shard = kwargs["shards"] + space_type = kwargs.get("space_type") + if space_type is None: + if method_name == "hnsw": + space_type = "l2" + else: + space_type = "cosine" + data_type = kwargs["data_type"] + vector_field = kwargs.get("vector_field", Field.VECTOR.value) + using_ugc = kwargs.get("using_ugc", False) + + if method_name == "ivfpq": + ivfpq_m = kwargs["ivfpq_m"] + nlist = kwargs["nlist"] + centroids_use_hnsw = nlist > 10000 + centroids_hnsw_m = 24 + centroids_hnsw_ef_construct = 500 + centroids_hnsw_ef_search = 100 + parameters = { + "m": ivfpq_m, + "nlist": nlist, + "centroids_use_hnsw": centroids_use_hnsw, + "centroids_hnsw_m": centroids_hnsw_m, + "centroids_hnsw_ef_construct": centroids_hnsw_ef_construct, + "centroids_hnsw_ef_search": centroids_hnsw_ef_search, + } + elif method_name == "hnsw": + neighbor = kwargs["hnsw_m"] + ef_construction = kwargs["hnsw_ef_construction"] + parameters = {"m": neighbor, "ef_construction": ef_construction} + elif method_name == "flat": + parameters = {} + else: + raise RuntimeError(f"unexpected method_name: {method_name}") + + mapping = { + "settings": {"index": {"number_of_shards": shard, "knn": True}}, + "mappings": { + "properties": { + vector_field: { + "type": "knn_vector", + "dimension": dimension, + "data_type": data_type, + "method": { + "engine": engine, + "name": method_name, + "space_type": space_type, + "parameters": parameters, + }, + }, + text_field: {"type": "text", "analyzer": analyzer}, + } + }, + } + + if excludes_from_source: + mapping["mappings"]["_source"] = {"excludes": excludes_from_source} # e.g. {"excludes": ["vector_field"]} + + if using_ugc and method_name == "ivfpq": + mapping["settings"]["index"]["knn_routing"] = True + mapping["settings"]["index"]["knn.offline.construction"] = True + elif using_ugc and method_name == "hnsw" or using_ugc and method_name == "flat": + mapping["settings"]["index"]["knn_routing"] = True + return mapping + + +def default_text_search_query( + query_text: str, + k: int = 4, + text_field: str = Field.CONTENT_KEY.value, + must: Optional[list[dict]] = None, + must_not: Optional[list[dict]] = None, + should: Optional[list[dict]] = None, + minimum_should_match: int = 0, + filters: Optional[list[dict]] = None, + routing: Optional[str] = None, + routing_field: Optional[str] = None, + **kwargs, +) -> dict: + query_clause: dict[str, Any] = {} + if routing is not None: + query_clause = { + "bool": {"must": [{"match": {text_field: query_text}}, {"term": {f"{routing_field}.keyword": routing}}]} + } + else: + query_clause = {"match": {text_field: query_text}} + # build the simplest search_query when only query_text is specified + if not must and not must_not and not should and not filters: + search_query = {"size": k, "query": query_clause} + return search_query + + # build complex search_query when either of must/must_not/should/filter is specified + if must: + if not isinstance(must, list): + raise RuntimeError(f"unexpected [must] clause with {type(filters)}") + if query_clause not in must: + must.append(query_clause) + else: + must = [query_clause] + + boolean_query: dict[str, Any] = {"must": must} + + if must_not: + if not isinstance(must_not, list): + raise RuntimeError(f"unexpected [must_not] clause with {type(filters)}") + boolean_query["must_not"] = must_not + + if should: + if not isinstance(should, list): + raise RuntimeError(f"unexpected [should] clause with {type(filters)}") + boolean_query["should"] = should + if minimum_should_match != 0: + boolean_query["minimum_should_match"] = minimum_should_match + + if filters: + if not isinstance(filters, list): + raise RuntimeError(f"unexpected [filter] clause with {type(filters)}") + boolean_query["filter"] = filters + + search_query = {"size": k, "query": {"bool": boolean_query}} + return search_query + + +def default_vector_search_query( + query_vector: list[float], + k: int = 4, + min_score: str = "0.0", + ef_search: Optional[str] = None, # only for hnsw + nprobe: Optional[str] = None, # "2000" + reorder_factor: Optional[str] = None, # "20" + client_refactor: Optional[str] = None, # "true" + vector_field: str = Field.VECTOR.value, + filters: Optional[list[dict]] = None, + filter_type: Optional[str] = None, + **kwargs, +) -> dict: + if filters is not None: + filter_type = "post_filter" if filter_type is None else filter_type + if not isinstance(filters, list): + raise RuntimeError(f"unexpected filter with {type(filters)}") + final_ext: dict[str, Any] = {"lvector": {}} + if min_score != "0.0": + final_ext["lvector"]["min_score"] = min_score + if ef_search: + final_ext["lvector"]["ef_search"] = ef_search + if nprobe: + final_ext["lvector"]["nprobe"] = nprobe + if reorder_factor: + final_ext["lvector"]["reorder_factor"] = reorder_factor + if client_refactor: + final_ext["lvector"]["client_refactor"] = client_refactor + + search_query: dict[str, Any] = { + "size": k, + "_source": True, # force return '_source' + "query": {"knn": {vector_field: {"vector": query_vector, "k": k}}}, + } + + if filters is not None: + # when using filter, transform filter from List[Dict] to Dict as valid format + filter_dict = {"bool": {"must": filters}} if len(filters) > 1 else filters[0] + search_query["query"]["knn"][vector_field]["filter"] = filter_dict # filter should be Dict + if filter_type: + final_ext["lvector"]["filter_type"] = filter_type + + if final_ext != {"lvector": {}}: + search_query["ext"] = final_ext + return search_query + + +class LindormVectorStoreFactory(AbstractVectorFactory): + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> LindormVectorStore: + lindorm_config = LindormVectorStoreConfig( + hosts=dify_config.LINDORM_URL or "", + username=dify_config.LINDORM_USERNAME, + password=dify_config.LINDORM_PASSWORD, + using_ugc=dify_config.USING_UGC_INDEX, + ) + using_ugc = dify_config.USING_UGC_INDEX + if using_ugc is None: + raise ValueError("USING_UGC_INDEX is not set") + routing_value = None + if dataset.index_struct: + # if an existed record's index_struct_dict doesn't contain using_ugc field, + # it actually stores in the normal index format + stored_in_ugc: bool = dataset.index_struct_dict.get("using_ugc", False) + using_ugc = stored_in_ugc + if stored_in_ugc: + dimension = dataset.index_struct_dict["dimension"] + index_type = dataset.index_struct_dict["index_type"] + distance_type = dataset.index_struct_dict["distance_type"] + routing_value = dataset.index_struct_dict["vector_store"]["class_prefix"] + index_name = f"{UGC_INDEX_PREFIX}_{dimension}_{index_type}_{distance_type}" + else: + index_name = dataset.index_struct_dict["vector_store"]["class_prefix"] + else: + embedding_vector = embeddings.embed_query("hello word") + dimension = len(embedding_vector) + index_type = dify_config.DEFAULT_INDEX_TYPE + distance_type = dify_config.DEFAULT_DISTANCE_TYPE + class_prefix = Dataset.gen_collection_name_by_id(dataset.id) + index_struct_dict = { + "type": VectorType.LINDORM, + "vector_store": {"class_prefix": class_prefix}, + "index_type": index_type, + "dimension": dimension, + "distance_type": distance_type, + "using_ugc": using_ugc, + } + dataset.index_struct = json.dumps(index_struct_dict) + if using_ugc: + index_name = f"{UGC_INDEX_PREFIX}_{dimension}_{index_type}_{distance_type}" + routing_value = class_prefix + else: + index_name = class_prefix + return LindormVectorStore(index_name, lindorm_config, routing_value=routing_value, using_ugc=using_ugc) diff --git a/api/core/rag/datasource/vdb/milvus/__init__.py b/api/core/rag/datasource/vdb/milvus/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/milvus/milvus_vector.py b/api/core/rag/datasource/vdb/milvus/milvus_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..9a184f7dd99ad95e610703168d2dc3a3d9ac7284 --- /dev/null +++ b/api/core/rag/datasource/vdb/milvus/milvus_vector.py @@ -0,0 +1,368 @@ +import json +import logging +from typing import Any, Optional + +from packaging import version +from pydantic import BaseModel, model_validator +from pymilvus import MilvusClient, MilvusException # type: ignore +from pymilvus.milvus_client import IndexParams # type: ignore + +from configs import dify_config +from core.rag.datasource.vdb.field import Field +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from extensions.ext_redis import redis_client +from models.dataset import Dataset + +logger = logging.getLogger(__name__) + + +class MilvusConfig(BaseModel): + """ + Configuration class for Milvus connection. + """ + + uri: str # Milvus server URI + token: Optional[str] = None # Optional token for authentication + user: str # Username for authentication + password: str # Password for authentication + batch_size: int = 100 # Batch size for operations + database: str = "default" # Database name + enable_hybrid_search: bool = False # Flag to enable hybrid search + + @model_validator(mode="before") + @classmethod + def validate_config(cls, values: dict) -> dict: + """ + Validate the configuration values. + Raises ValueError if required fields are missing. + """ + if not values.get("uri"): + raise ValueError("config MILVUS_URI is required") + if not values.get("user"): + raise ValueError("config MILVUS_USER is required") + if not values.get("password"): + raise ValueError("config MILVUS_PASSWORD is required") + return values + + def to_milvus_params(self): + """ + Convert the configuration to a dictionary of Milvus connection parameters. + """ + return { + "uri": self.uri, + "token": self.token, + "user": self.user, + "password": self.password, + "db_name": self.database, + } + + +class MilvusVector(BaseVector): + """ + Milvus vector storage implementation. + """ + + def __init__(self, collection_name: str, config: MilvusConfig): + super().__init__(collection_name) + self._client_config = config + self._client = self._init_client(config) + self._consistency_level = "Session" # Consistency level for Milvus operations + self._fields: list[str] = [] # List of fields in the collection + self._hybrid_search_enabled = self._check_hybrid_search_support() # Check if hybrid search is supported + + def _check_hybrid_search_support(self) -> bool: + """ + Check if the current Milvus version supports hybrid search. + Returns True if the version is >= 2.5.0, otherwise False. + """ + if not self._client_config.enable_hybrid_search: + return False + + try: + milvus_version = self._client.get_server_version() + return version.parse(milvus_version).base_version >= version.parse("2.5.0").base_version + except Exception as e: + logger.warning(f"Failed to check Milvus version: {str(e)}. Disabling hybrid search.") + return False + + def get_type(self) -> str: + """ + Get the type of vector storage (Milvus). + """ + return VectorType.MILVUS + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + """ + Create a collection and add texts with embeddings. + """ + index_params = {"metric_type": "IP", "index_type": "HNSW", "params": {"M": 8, "efConstruction": 64}} + metadatas = [d.metadata if d.metadata is not None else {} for d in texts] + self.create_collection(embeddings, metadatas, index_params) + self.add_texts(texts, embeddings) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + """ + Add texts and their embeddings to the collection. + """ + insert_dict_list = [] + for i in range(len(documents)): + insert_dict = { + # Do not need to insert the sparse_vector field separately, as the text_bm25_emb + # function will automatically convert the native text into a sparse vector for us. + Field.CONTENT_KEY.value: documents[i].page_content, + Field.VECTOR.value: embeddings[i], + Field.METADATA_KEY.value: documents[i].metadata, + } + insert_dict_list.append(insert_dict) + # Total insert count + total_count = len(insert_dict_list) + pks: list[str] = [] + + for i in range(0, total_count, 1000): + # Insert into the collection. + batch_insert_list = insert_dict_list[i : i + 1000] + try: + ids = self._client.insert(collection_name=self._collection_name, data=batch_insert_list) + pks.extend(ids) + except MilvusException as e: + logger.exception("Failed to insert batch starting at entity: %s/%s", i, total_count) + raise e + return pks + + def get_ids_by_metadata_field(self, key: str, value: str): + """ + Get document IDs by metadata field key and value. + """ + result = self._client.query( + collection_name=self._collection_name, filter=f'metadata["{key}"] == "{value}"', output_fields=["id"] + ) + if result: + return [item["id"] for item in result] + else: + return None + + def delete_by_metadata_field(self, key: str, value: str): + """ + Delete documents by metadata field key and value. + """ + if self._client.has_collection(self._collection_name): + ids = self.get_ids_by_metadata_field(key, value) + if ids: + self._client.delete(collection_name=self._collection_name, pks=ids) + + def delete_by_ids(self, ids: list[str]) -> None: + """ + Delete documents by their IDs. + """ + if self._client.has_collection(self._collection_name): + result = self._client.query( + collection_name=self._collection_name, filter=f'metadata["doc_id"] in {ids}', output_fields=["id"] + ) + if result: + ids = [item["id"] for item in result] + self._client.delete(collection_name=self._collection_name, pks=ids) + + def delete(self) -> None: + """ + Delete the entire collection. + """ + if self._client.has_collection(self._collection_name): + self._client.drop_collection(self._collection_name, None) + + def text_exists(self, id: str) -> bool: + """ + Check if a text with the given ID exists in the collection. + """ + if not self._client.has_collection(self._collection_name): + return False + + result = self._client.query( + collection_name=self._collection_name, filter=f'metadata["doc_id"] == "{id}"', output_fields=["id"] + ) + + return len(result) > 0 + + def field_exists(self, field: str) -> bool: + """ + Check if a field exists in the collection. + """ + return field in self._fields + + def _process_search_results( + self, results: list[Any], output_fields: list[str], score_threshold: float = 0.0 + ) -> list[Document]: + """ + Common method to process search results + + :param results: Search results + :param output_fields: Fields to be output + :param score_threshold: Score threshold for filtering + :return: List of documents + """ + docs = [] + for result in results[0]: + metadata = result["entity"].get(output_fields[1], {}) + metadata["score"] = result["distance"] + + if result["distance"] > score_threshold: + doc = Document(page_content=result["entity"].get(output_fields[0], ""), metadata=metadata) + docs.append(doc) + + return docs + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + """ + Search for documents by vector similarity. + """ + results = self._client.search( + collection_name=self._collection_name, + data=[query_vector], + anns_field=Field.VECTOR.value, + limit=kwargs.get("top_k", 4), + output_fields=[Field.CONTENT_KEY.value, Field.METADATA_KEY.value], + ) + + return self._process_search_results( + results, + output_fields=[Field.CONTENT_KEY.value, Field.METADATA_KEY.value], + score_threshold=float(kwargs.get("score_threshold") or 0.0), + ) + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + """ + Search for documents by full-text search (if hybrid search is enabled). + """ + if not self._hybrid_search_enabled or not self.field_exists(Field.SPARSE_VECTOR.value): + logger.warning("Full-text search is not supported in current Milvus version (requires >= 2.5.0)") + return [] + + results = self._client.search( + collection_name=self._collection_name, + data=[query], + anns_field=Field.SPARSE_VECTOR.value, + limit=kwargs.get("top_k", 4), + output_fields=[Field.CONTENT_KEY.value, Field.METADATA_KEY.value], + ) + + return self._process_search_results( + results, + output_fields=[Field.CONTENT_KEY.value, Field.METADATA_KEY.value], + score_threshold=float(kwargs.get("score_threshold") or 0.0), + ) + + def create_collection( + self, embeddings: list, metadatas: Optional[list[dict]] = None, index_params: Optional[dict] = None + ): + """ + Create a new collection in Milvus with the specified schema and index parameters. + """ + lock_name = "vector_indexing_lock_{}".format(self._collection_name) + with redis_client.lock(lock_name, timeout=20): + collection_exist_cache_key = "vector_indexing_{}".format(self._collection_name) + if redis_client.get(collection_exist_cache_key): + return + # Grab the existing collection if it exists + if not self._client.has_collection(self._collection_name): + from pymilvus import CollectionSchema, DataType, FieldSchema, Function, FunctionType # type: ignore + from pymilvus.orm.types import infer_dtype_bydata # type: ignore + + # Determine embedding dim + dim = len(embeddings[0]) + fields = [] + if metadatas: + fields.append(FieldSchema(Field.METADATA_KEY.value, DataType.JSON, max_length=65_535)) + + # Create the text field, enable_analyzer will be set True to support milvus automatically + # transfer text to sparse_vector, reference: https://milvus.io/docs/full-text-search.md + fields.append( + FieldSchema( + Field.CONTENT_KEY.value, + DataType.VARCHAR, + max_length=65_535, + enable_analyzer=self._hybrid_search_enabled, + ) + ) + # Create the primary key field + fields.append(FieldSchema(Field.PRIMARY_KEY.value, DataType.INT64, is_primary=True, auto_id=True)) + # Create the vector field, supports binary or float vectors + fields.append(FieldSchema(Field.VECTOR.value, infer_dtype_bydata(embeddings[0]), dim=dim)) + # Create Sparse Vector Index for the collection + if self._hybrid_search_enabled: + fields.append(FieldSchema(Field.SPARSE_VECTOR.value, DataType.SPARSE_FLOAT_VECTOR)) + + schema = CollectionSchema(fields) + + # Create custom function to support text to sparse vector by BM25 + if self._hybrid_search_enabled: + bm25_function = Function( + name="text_bm25_emb", + input_field_names=[Field.CONTENT_KEY.value], + output_field_names=[Field.SPARSE_VECTOR.value], + function_type=FunctionType.BM25, + ) + schema.add_function(bm25_function) + + for x in schema.fields: + self._fields.append(x.name) + # Since primary field is auto-id, no need to track it + self._fields.remove(Field.PRIMARY_KEY.value) + + # Create Index params for the collection + index_params_obj = IndexParams() + index_params_obj.add_index(field_name=Field.VECTOR.value, **index_params) + + # Create Sparse Vector Index for the collection + if self._hybrid_search_enabled: + index_params_obj.add_index( + field_name=Field.SPARSE_VECTOR.value, index_type="AUTOINDEX", metric_type="BM25" + ) + + # Create the collection + self._client.create_collection( + collection_name=self._collection_name, + schema=schema, + index_params=index_params_obj, + consistency_level=self._consistency_level, + ) + redis_client.set(collection_exist_cache_key, 1, ex=3600) + + def _init_client(self, config) -> MilvusClient: + """ + Initialize and return a Milvus client. + """ + client = MilvusClient(uri=config.uri, user=config.user, password=config.password, db_name=config.database) + return client + + +class MilvusVectorFactory(AbstractVectorFactory): + """ + Factory class for creating MilvusVector instances. + """ + + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> MilvusVector: + """ + Initialize a MilvusVector instance for the given dataset. + """ + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id) + dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.MILVUS, collection_name)) + + return MilvusVector( + collection_name=collection_name, + config=MilvusConfig( + uri=dify_config.MILVUS_URI or "", + token=dify_config.MILVUS_TOKEN or "", + user=dify_config.MILVUS_USER or "", + password=dify_config.MILVUS_PASSWORD or "", + database=dify_config.MILVUS_DATABASE or "", + enable_hybrid_search=dify_config.MILVUS_ENABLE_HYBRID_SEARCH or False, + ), + ) diff --git a/api/core/rag/datasource/vdb/myscale/__init__.py b/api/core/rag/datasource/vdb/myscale/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/myscale/myscale_vector.py b/api/core/rag/datasource/vdb/myscale/myscale_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..556b952ec262b7515c2d7da7aa6a98ea354e39c6 --- /dev/null +++ b/api/core/rag/datasource/vdb/myscale/myscale_vector.py @@ -0,0 +1,175 @@ +import json +import logging +import uuid +from enum import Enum +from typing import Any + +from clickhouse_connect import get_client +from pydantic import BaseModel + +from configs import dify_config +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from models.dataset import Dataset + + +class MyScaleConfig(BaseModel): + host: str + port: int + user: str + password: str + database: str + fts_params: str + + +class SortOrder(Enum): + ASC = "ASC" + DESC = "DESC" + + +class MyScaleVector(BaseVector): + def __init__(self, collection_name: str, config: MyScaleConfig, metric: str = "Cosine"): + super().__init__(collection_name) + self._config = config + self._metric = metric + self._vec_order = SortOrder.ASC if metric.upper() in {"COSINE", "L2"} else SortOrder.DESC + self._client = get_client( + host=config.host, + port=config.port, + username=config.user, + password=config.password, + ) + self._client.command("SET allow_experimental_object_type=1") + + def get_type(self) -> str: + return VectorType.MYSCALE + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + dimension = len(embeddings[0]) + self._create_collection(dimension) + return self.add_texts(documents=texts, embeddings=embeddings, **kwargs) + + def _create_collection(self, dimension: int): + logging.info(f"create MyScale collection {self._collection_name} with dimension {dimension}") + self._client.command(f"CREATE DATABASE IF NOT EXISTS {self._config.database}") + fts_params = f"('{self._config.fts_params}')" if self._config.fts_params else "" + sql = f""" + CREATE TABLE IF NOT EXISTS {self._config.database}.{self._collection_name}( + id String, + text String, + vector Array(Float32), + metadata JSON, + CONSTRAINT cons_vec_len CHECK length(vector) = {dimension}, + VECTOR INDEX vidx vector TYPE DEFAULT('metric_type = {self._metric}'), + INDEX text_idx text TYPE fts{fts_params} + ) ENGINE = MergeTree ORDER BY id + """ + self._client.command(sql) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + ids = [] + columns = ["id", "text", "vector", "metadata"] + values = [] + for i, doc in enumerate(documents): + if doc.metadata is not None: + doc_id = doc.metadata.get("doc_id", str(uuid.uuid4())) + row = ( + doc_id, + self.escape_str(doc.page_content), + embeddings[i], + json.dumps(doc.metadata) if doc.metadata else {}, + ) + values.append(str(row)) + ids.append(doc_id) + sql = f""" + INSERT INTO {self._config.database}.{self._collection_name} + ({",".join(columns)}) VALUES {",".join(values)} + """ + self._client.command(sql) + return ids + + @staticmethod + def escape_str(value: Any) -> str: + return "".join(" " if c in {"\\", "'"} else c for c in str(value)) + + def text_exists(self, id: str) -> bool: + results = self._client.query(f"SELECT id FROM {self._config.database}.{self._collection_name} WHERE id='{id}'") + return results.row_count > 0 + + def delete_by_ids(self, ids: list[str]) -> None: + if not ids: + return + self._client.command( + f"DELETE FROM {self._config.database}.{self._collection_name} WHERE id IN {str(tuple(ids))}" + ) + + def get_ids_by_metadata_field(self, key: str, value: str): + rows = self._client.query( + f"SELECT DISTINCT id FROM {self._config.database}.{self._collection_name} WHERE metadata.{key}='{value}'" + ).result_rows + return [row[0] for row in rows] + + def delete_by_metadata_field(self, key: str, value: str) -> None: + self._client.command( + f"DELETE FROM {self._config.database}.{self._collection_name} WHERE metadata.{key}='{value}'" + ) + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + return self._search(f"distance(vector, {str(query_vector)})", self._vec_order, **kwargs) + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + return self._search(f"TextSearch('enable_nlq=false')(text, '{query}')", SortOrder.DESC, **kwargs) + + def _search(self, dist: str, order: SortOrder, **kwargs: Any) -> list[Document]: + top_k = kwargs.get("top_k", 4) + score_threshold = float(kwargs.get("score_threshold") or 0.0) + where_str = ( + f"WHERE dist < {1 - score_threshold}" + if self._metric.upper() == "COSINE" and order == SortOrder.ASC and score_threshold > 0.0 + else "" + ) + sql = f""" + SELECT text, vector, metadata, {dist} as dist FROM {self._config.database}.{self._collection_name} + {where_str} ORDER BY dist {order.value} LIMIT {top_k} + """ + try: + return [ + Document( + page_content=r["text"], + vector=r["vector"], + metadata=r["metadata"], + ) + for r in self._client.query(sql).named_results() + ] + except Exception as e: + logging.exception(f"\033[91m\033[1m{type(e)}\033[0m \033[95m{str(e)}\033[0m") # noqa:TRY401 + return [] + + def delete(self) -> None: + self._client.command(f"DROP TABLE IF EXISTS {self._config.database}.{self._collection_name}") + + +class MyScaleVectorFactory(AbstractVectorFactory): + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> MyScaleVector: + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix.lower() + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id).lower() + dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.MYSCALE, collection_name)) + + return MyScaleVector( + collection_name=collection_name, + config=MyScaleConfig( + host=dify_config.MYSCALE_HOST, + port=dify_config.MYSCALE_PORT, + user=dify_config.MYSCALE_USER, + password=dify_config.MYSCALE_PASSWORD, + database=dify_config.MYSCALE_DATABASE, + fts_params=dify_config.MYSCALE_FTS_PARAMS, + ), + ) diff --git a/api/core/rag/datasource/vdb/oceanbase/__init__.py b/api/core/rag/datasource/vdb/oceanbase/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/oceanbase/oceanbase_vector.py b/api/core/rag/datasource/vdb/oceanbase/oceanbase_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..3c2d53ce780ed82f0023468024df91fabea9f246 --- /dev/null +++ b/api/core/rag/datasource/vdb/oceanbase/oceanbase_vector.py @@ -0,0 +1,210 @@ +import json +import logging +import math +from typing import Any + +from pydantic import BaseModel, model_validator +from pyobvector import VECTOR, ObVecClient # type: ignore +from sqlalchemy import JSON, Column, String, func +from sqlalchemy.dialects.mysql import LONGTEXT + +from configs import dify_config +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from extensions.ext_redis import redis_client +from models.dataset import Dataset + +logger = logging.getLogger(__name__) + +DEFAULT_OCEANBASE_HNSW_BUILD_PARAM = {"M": 16, "efConstruction": 256} +DEFAULT_OCEANBASE_HNSW_SEARCH_PARAM = {"efSearch": 64} +OCEANBASE_SUPPORTED_VECTOR_INDEX_TYPE = "HNSW" +DEFAULT_OCEANBASE_VECTOR_METRIC_TYPE = "l2" + + +class OceanBaseVectorConfig(BaseModel): + host: str + port: int + user: str + password: str + database: str + + @model_validator(mode="before") + @classmethod + def validate_config(cls, values: dict) -> dict: + if not values["host"]: + raise ValueError("config OCEANBASE_VECTOR_HOST is required") + if not values["port"]: + raise ValueError("config OCEANBASE_VECTOR_PORT is required") + if not values["user"]: + raise ValueError("config OCEANBASE_VECTOR_USER is required") + if not values["database"]: + raise ValueError("config OCEANBASE_VECTOR_DATABASE is required") + return values + + +class OceanBaseVector(BaseVector): + def __init__(self, collection_name: str, config: OceanBaseVectorConfig): + super().__init__(collection_name) + self._config = config + self._hnsw_ef_search = -1 + self._client = ObVecClient( + uri=f"{self._config.host}:{self._config.port}", + user=self._config.user, + password=self._config.password, + db_name=self._config.database, + ) + + def get_type(self) -> str: + return VectorType.OCEANBASE + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + self._vec_dim = len(embeddings[0]) + self._create_collection() + self.add_texts(texts, embeddings) + + def _create_collection(self) -> None: + lock_name = "vector_indexing_lock_" + self._collection_name + with redis_client.lock(lock_name, timeout=20): + collection_exist_cache_key = "vector_indexing_" + self._collection_name + if redis_client.get(collection_exist_cache_key): + return + + if self._client.check_table_exists(self._collection_name): + return + + self.delete() + + cols = [ + Column("id", String(36), primary_key=True, autoincrement=False), + Column("vector", VECTOR(self._vec_dim)), + Column("text", LONGTEXT), + Column("metadata", JSON), + ] + vidx_params = self._client.prepare_index_params() + vidx_params.add_index( + field_name="vector", + index_type=OCEANBASE_SUPPORTED_VECTOR_INDEX_TYPE, + index_name="vector_index", + metric_type=DEFAULT_OCEANBASE_VECTOR_METRIC_TYPE, + params=DEFAULT_OCEANBASE_HNSW_BUILD_PARAM, + ) + + self._client.create_table_with_index_params( + table_name=self._collection_name, + columns=cols, + vidxs=vidx_params, + ) + vals = [] + params = self._client.perform_raw_text_sql("SHOW PARAMETERS LIKE '%ob_vector_memory_limit_percentage%'") + for row in params: + val = int(row[6]) + vals.append(val) + if len(vals) == 0: + raise ValueError("ob_vector_memory_limit_percentage not found in parameters.") + if any(val == 0 for val in vals): + try: + self._client.perform_raw_text_sql("ALTER SYSTEM SET ob_vector_memory_limit_percentage = 30") + except Exception as e: + raise Exception( + "Failed to set ob_vector_memory_limit_percentage. " + + "Maybe the database user has insufficient privilege.", + e, + ) + redis_client.set(collection_exist_cache_key, 1, ex=3600) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + ids = self._get_uuids(documents) + for id, doc, emb in zip(ids, documents, embeddings): + self._client.insert( + table_name=self._collection_name, + data={ + "id": id, + "vector": emb, + "text": doc.page_content, + "metadata": doc.metadata, + }, + ) + + def text_exists(self, id: str) -> bool: + cur = self._client.get(table_name=self._collection_name, id=id) + return bool(cur.rowcount != 0) + + def delete_by_ids(self, ids: list[str]) -> None: + if not ids: + return + self._client.delete(table_name=self._collection_name, ids=ids) + + def get_ids_by_metadata_field(self, key: str, value: str) -> list[str]: + cur = self._client.get( + table_name=self._collection_name, + where_clause=f"metadata->>'$.{key}' = '{value}'", + output_column_name=["id"], + ) + return [row[0] for row in cur] + + def delete_by_metadata_field(self, key: str, value: str) -> None: + ids = self.get_ids_by_metadata_field(key, value) + self.delete_by_ids(ids) + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + return [] + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + ef_search = kwargs.get("ef_search", self._hnsw_ef_search) + if ef_search != self._hnsw_ef_search: + self._client.set_ob_hnsw_ef_search(ef_search) + self._hnsw_ef_search = ef_search + topk = kwargs.get("top_k", 10) + cur = self._client.ann_search( + table_name=self._collection_name, + vec_column_name="vector", + vec_data=query_vector, + topk=topk, + distance_func=func.l2_distance, + output_column_names=["text", "metadata"], + with_dist=True, + ) + docs = [] + for text, metadata, distance in cur: + metadata = json.loads(metadata) + metadata["score"] = 1 - distance / math.sqrt(2) + docs.append( + Document( + page_content=text, + metadata=metadata, + ) + ) + return docs + + def delete(self) -> None: + self._client.drop_table_if_exist(self._collection_name) + + +class OceanBaseVectorFactory(AbstractVectorFactory): + def init_vector( + self, + dataset: Dataset, + attributes: list, + embeddings: Embeddings, + ) -> BaseVector: + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix.lower() + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id).lower() + dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.OCEANBASE, collection_name)) + return OceanBaseVector( + collection_name, + OceanBaseVectorConfig( + host=dify_config.OCEANBASE_VECTOR_HOST or "", + port=dify_config.OCEANBASE_VECTOR_PORT or 0, + user=dify_config.OCEANBASE_VECTOR_USER or "", + password=(dify_config.OCEANBASE_VECTOR_PASSWORD or ""), + database=dify_config.OCEANBASE_VECTOR_DATABASE or "", + ), + ) diff --git a/api/core/rag/datasource/vdb/opensearch/__init__.py b/api/core/rag/datasource/vdb/opensearch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/opensearch/opensearch_vector.py b/api/core/rag/datasource/vdb/opensearch/opensearch_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..72a15022052f0aefb6bc0833b8c1ddcf87b5efd1 --- /dev/null +++ b/api/core/rag/datasource/vdb/opensearch/opensearch_vector.py @@ -0,0 +1,254 @@ +import json +import logging +import ssl +from typing import Any, Optional +from uuid import uuid4 + +from opensearchpy import OpenSearch, helpers +from opensearchpy.helpers import BulkIndexError +from pydantic import BaseModel, model_validator + +from configs import dify_config +from core.rag.datasource.vdb.field import Field +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from extensions.ext_redis import redis_client +from models.dataset import Dataset + +logger = logging.getLogger(__name__) + + +class OpenSearchConfig(BaseModel): + host: str + port: int + user: Optional[str] = None + password: Optional[str] = None + secure: bool = False + + @model_validator(mode="before") + @classmethod + def validate_config(cls, values: dict) -> dict: + if not values.get("host"): + raise ValueError("config OPENSEARCH_HOST is required") + if not values.get("port"): + raise ValueError("config OPENSEARCH_PORT is required") + return values + + def create_ssl_context(self) -> ssl.SSLContext: + ssl_context = ssl.create_default_context() + ssl_context.check_hostname = False + ssl_context.verify_mode = ssl.CERT_NONE # Disable Certificate Validation + return ssl_context + + def to_opensearch_params(self) -> dict[str, Any]: + params = { + "hosts": [{"host": self.host, "port": self.port}], + "use_ssl": self.secure, + "verify_certs": self.secure, + } + if self.user and self.password: + params["http_auth"] = (self.user, self.password) + if self.secure: + params["ssl_context"] = self.create_ssl_context() + return params + + +class OpenSearchVector(BaseVector): + def __init__(self, collection_name: str, config: OpenSearchConfig): + super().__init__(collection_name) + self._client_config = config + self._client = OpenSearch(**config.to_opensearch_params()) + + def get_type(self) -> str: + return VectorType.OPENSEARCH + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + metadatas = [d.metadata if d.metadata is not None else {} for d in texts] + self.create_collection(embeddings, metadatas) + self.add_texts(texts, embeddings) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + actions = [] + for i in range(len(documents)): + action = { + "_op_type": "index", + "_index": self._collection_name.lower(), + "_id": uuid4().hex, + "_source": { + Field.CONTENT_KEY.value: documents[i].page_content, + Field.VECTOR.value: embeddings[i], # Make sure you pass an array here + Field.METADATA_KEY.value: documents[i].metadata, + }, + } + actions.append(action) + + helpers.bulk(self._client, actions) + + def get_ids_by_metadata_field(self, key: str, value: str): + query = {"query": {"term": {f"{Field.METADATA_KEY.value}.{key}": value}}} + response = self._client.search(index=self._collection_name.lower(), body=query) + if response["hits"]["hits"]: + return [hit["_id"] for hit in response["hits"]["hits"]] + else: + return None + + def delete_by_metadata_field(self, key: str, value: str): + ids = self.get_ids_by_metadata_field(key, value) + if ids: + self.delete_by_ids(ids) + + def delete_by_ids(self, ids: list[str]) -> None: + index_name = self._collection_name.lower() + if not self._client.indices.exists(index=index_name): + logger.warning(f"Index {index_name} does not exist") + return + + # Obtaining All Actual Documents_ID + actual_ids = [] + + for doc_id in ids: + es_ids = self.get_ids_by_metadata_field("doc_id", doc_id) + if es_ids: + actual_ids.extend(es_ids) + else: + logger.warning(f"Document with metadata doc_id {doc_id} not found for deletion") + + if actual_ids: + actions = [{"_op_type": "delete", "_index": index_name, "_id": es_id} for es_id in actual_ids] + try: + helpers.bulk(self._client, actions) + except BulkIndexError as e: + for error in e.errors: + delete_error = error.get("delete", {}) + status = delete_error.get("status") + doc_id = delete_error.get("_id") + + if status == 404: + logger.warning(f"Document not found for deletion: {doc_id}") + else: + logger.exception(f"Error deleting document: {error}") + + def delete(self) -> None: + self._client.indices.delete(index=self._collection_name.lower()) + + def text_exists(self, id: str) -> bool: + try: + self._client.get(index=self._collection_name.lower(), id=id) + return True + except: + return False + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + # Make sure query_vector is a list + if not isinstance(query_vector, list): + raise ValueError("query_vector should be a list of floats") + + # Check whether query_vector is a floating-point number list + if not all(isinstance(x, float) for x in query_vector): + raise ValueError("All elements in query_vector should be floats") + + query = { + "size": kwargs.get("top_k", 4), + "query": {"knn": {Field.VECTOR.value: {Field.VECTOR.value: query_vector, "k": kwargs.get("top_k", 4)}}}, + } + + try: + response = self._client.search(index=self._collection_name.lower(), body=query) + except Exception as e: + logger.exception(f"Error executing vector search, query: {query}") + raise + + docs = [] + for hit in response["hits"]["hits"]: + metadata = hit["_source"].get(Field.METADATA_KEY.value, {}) + + # Make sure metadata is a dictionary + if metadata is None: + metadata = {} + + metadata["score"] = hit["_score"] + score_threshold = float(kwargs.get("score_threshold") or 0.0) + if hit["_score"] > score_threshold: + doc = Document(page_content=hit["_source"].get(Field.CONTENT_KEY.value), metadata=metadata) + docs.append(doc) + + return docs + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + full_text_query = {"query": {"match": {Field.CONTENT_KEY.value: query}}} + + response = self._client.search(index=self._collection_name.lower(), body=full_text_query) + + docs = [] + for hit in response["hits"]["hits"]: + metadata = hit["_source"].get(Field.METADATA_KEY.value) + vector = hit["_source"].get(Field.VECTOR.value) + page_content = hit["_source"].get(Field.CONTENT_KEY.value) + doc = Document(page_content=page_content, vector=vector, metadata=metadata) + docs.append(doc) + + return docs + + def create_collection( + self, embeddings: list, metadatas: Optional[list[dict]] = None, index_params: Optional[dict] = None + ): + lock_name = f"vector_indexing_lock_{self._collection_name.lower()}" + with redis_client.lock(lock_name, timeout=20): + collection_exist_cache_key = f"vector_indexing_{self._collection_name.lower()}" + if redis_client.get(collection_exist_cache_key): + logger.info(f"Collection {self._collection_name.lower()} already exists.") + return + + if not self._client.indices.exists(index=self._collection_name.lower()): + index_body = { + "settings": {"index": {"knn": True}}, + "mappings": { + "properties": { + Field.CONTENT_KEY.value: {"type": "text"}, + Field.VECTOR.value: { + "type": "knn_vector", + "dimension": len(embeddings[0]), # Make sure the dimension is correct here + "method": { + "name": "hnsw", + "space_type": "l2", + "engine": "faiss", + "parameters": {"ef_construction": 64, "m": 8}, + }, + }, + Field.METADATA_KEY.value: { + "type": "object", + "properties": { + "doc_id": {"type": "keyword"} # Map doc_id to keyword type + }, + }, + } + }, + } + + self._client.indices.create(index=self._collection_name.lower(), body=index_body) + + redis_client.set(collection_exist_cache_key, 1, ex=3600) + + +class OpenSearchVectorFactory(AbstractVectorFactory): + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> OpenSearchVector: + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix.lower() + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id).lower() + dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.OPENSEARCH, collection_name)) + + open_search_config = OpenSearchConfig( + host=dify_config.OPENSEARCH_HOST or "localhost", + port=dify_config.OPENSEARCH_PORT, + user=dify_config.OPENSEARCH_USER, + password=dify_config.OPENSEARCH_PASSWORD, + secure=dify_config.OPENSEARCH_SECURE, + ) + + return OpenSearchVector(collection_name=collection_name, config=open_search_config) diff --git a/api/core/rag/datasource/vdb/oracle/__init__.py b/api/core/rag/datasource/vdb/oracle/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/oracle/oraclevector.py b/api/core/rag/datasource/vdb/oracle/oraclevector.py new file mode 100644 index 0000000000000000000000000000000000000000..a58df7eb9f403d427360d4d3feb4075b63d7e4d5 --- /dev/null +++ b/api/core/rag/datasource/vdb/oracle/oraclevector.py @@ -0,0 +1,296 @@ +import array +import json +import re +import uuid +from contextlib import contextmanager +from typing import Any + +import jieba.posseg as pseg # type: ignore +import numpy +import oracledb +from pydantic import BaseModel, model_validator + +from configs import dify_config +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from extensions.ext_redis import redis_client +from models.dataset import Dataset + +oracledb.defaults.fetch_lobs = False + + +class OracleVectorConfig(BaseModel): + host: str + port: int + user: str + password: str + database: str + + @model_validator(mode="before") + @classmethod + def validate_config(cls, values: dict) -> dict: + if not values["host"]: + raise ValueError("config ORACLE_HOST is required") + if not values["port"]: + raise ValueError("config ORACLE_PORT is required") + if not values["user"]: + raise ValueError("config ORACLE_USER is required") + if not values["password"]: + raise ValueError("config ORACLE_PASSWORD is required") + if not values["database"]: + raise ValueError("config ORACLE_DB is required") + return values + + +SQL_CREATE_TABLE = """ +CREATE TABLE IF NOT EXISTS {table_name} ( + id varchar2(100) + ,text CLOB NOT NULL + ,meta JSON + ,embedding vector NOT NULL +) +""" +SQL_CREATE_INDEX = """ +CREATE INDEX IF NOT EXISTS idx_docs_{table_name} ON {table_name}(text) +INDEXTYPE IS CTXSYS.CONTEXT PARAMETERS +('FILTER CTXSYS.NULL_FILTER SECTION GROUP CTXSYS.HTML_SECTION_GROUP LEXER sys.my_chinese_vgram_lexer') +""" + + +class OracleVector(BaseVector): + def __init__(self, collection_name: str, config: OracleVectorConfig): + super().__init__(collection_name) + self.pool = self._create_connection_pool(config) + self.table_name = f"embedding_{collection_name}" + + def get_type(self) -> str: + return VectorType.ORACLE + + def numpy_converter_in(self, value): + if value.dtype == numpy.float64: + dtype = "d" + elif value.dtype == numpy.float32: + dtype = "f" + else: + dtype = "b" + return array.array(dtype, value) + + def input_type_handler(self, cursor, value, arraysize): + if isinstance(value, numpy.ndarray): + return cursor.var( + oracledb.DB_TYPE_VECTOR, + arraysize=arraysize, + inconverter=self.numpy_converter_in, + ) + + def numpy_converter_out(self, value): + if value.typecode == "b": + return numpy.array(value, copy=False, dtype=numpy.int8) + elif value.typecode == "f": + return numpy.array(value, copy=False, dtype=numpy.float32) + else: + return numpy.array(value, copy=False, dtype=numpy.float64) + + def output_type_handler(self, cursor, metadata): + if metadata.type_code is oracledb.DB_TYPE_VECTOR: + return cursor.var( + metadata.type_code, + arraysize=cursor.arraysize, + outconverter=self.numpy_converter_out, + ) + + def _create_connection_pool(self, config: OracleVectorConfig): + return oracledb.create_pool( + user=config.user, + password=config.password, + dsn="{}:{}/{}".format(config.host, config.port, config.database), + min=1, + max=50, + increment=1, + ) + + @contextmanager + def _get_cursor(self): + conn = self.pool.acquire() + conn.inputtypehandler = self.input_type_handler + conn.outputtypehandler = self.output_type_handler + cur = conn.cursor() + try: + yield cur + finally: + cur.close() + conn.commit() + conn.close() + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + dimension = len(embeddings[0]) + self._create_collection(dimension) + return self.add_texts(texts, embeddings) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + values = [] + pks = [] + for i, doc in enumerate(documents): + if doc.metadata is not None: + doc_id = doc.metadata.get("doc_id", str(uuid.uuid4())) + pks.append(doc_id) + values.append( + ( + doc_id, + doc.page_content, + json.dumps(doc.metadata), + # array.array("f", embeddings[i]), + numpy.array(embeddings[i]), + ) + ) + # print(f"INSERT INTO {self.table_name} (id, text, meta, embedding) VALUES (:1, :2, :3, :4)") + with self._get_cursor() as cur: + cur.executemany( + f"INSERT INTO {self.table_name} (id, text, meta, embedding) VALUES (:1, :2, :3, :4)", values + ) + return pks + + def text_exists(self, id: str) -> bool: + with self._get_cursor() as cur: + cur.execute(f"SELECT id FROM {self.table_name} WHERE id = '%s'" % (id,)) + return cur.fetchone() is not None + + def get_by_ids(self, ids: list[str]) -> list[Document]: + with self._get_cursor() as cur: + cur.execute(f"SELECT meta, text FROM {self.table_name} WHERE id IN %s", (tuple(ids),)) + docs = [] + for record in cur: + docs.append(Document(page_content=record[1], metadata=record[0])) + return docs + + def delete_by_ids(self, ids: list[str]) -> None: + if not ids: + return + with self._get_cursor() as cur: + cur.execute(f"DELETE FROM {self.table_name} WHERE id IN %s" % (tuple(ids),)) + + def delete_by_metadata_field(self, key: str, value: str) -> None: + with self._get_cursor() as cur: + cur.execute(f"DELETE FROM {self.table_name} WHERE meta->>%s = %s", (key, value)) + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + """ + Search the nearest neighbors to a vector. + + :param query_vector: The input vector to search for similar items. + :param top_k: The number of nearest neighbors to return, default is 5. + :return: List of Documents that are nearest to the query vector. + """ + top_k = kwargs.get("top_k", 4) + with self._get_cursor() as cur: + cur.execute( + f"SELECT meta, text, vector_distance(embedding,:1) AS distance FROM {self.table_name}" + f" ORDER BY distance fetch first {top_k} rows only", + [numpy.array(query_vector)], + ) + docs = [] + score_threshold = float(kwargs.get("score_threshold") or 0.0) + for record in cur: + metadata, text, distance = record + score = 1 - distance + metadata["score"] = score + if score > score_threshold: + docs.append(Document(page_content=text, metadata=metadata)) + return docs + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + # lazy import + import nltk # type: ignore + from nltk.corpus import stopwords # type: ignore + + top_k = kwargs.get("top_k", 5) + # just not implement fetch by score_threshold now, may be later + score_threshold = float(kwargs.get("score_threshold") or 0.0) + if len(query) > 0: + # Check which language the query is in + zh_pattern = re.compile("[\u4e00-\u9fa5]+") + match = zh_pattern.search(query) + entities = [] + # match: query condition maybe is a chinese sentence, so using Jieba split,else using nltk split + if match: + words = pseg.cut(query) + current_entity = "" + for word, pos in words: + if pos in {"nr", "Ng", "eng", "nz", "n", "ORG", "v"}: # nr: 人名, ns: 地名, nt: 机构名 + current_entity += word + else: + if current_entity: + entities.append(current_entity) + current_entity = "" + if current_entity: + entities.append(current_entity) + else: + try: + nltk.data.find("tokenizers/punkt") + nltk.data.find("corpora/stopwords") + except LookupError: + nltk.download("punkt") + nltk.download("stopwords") + e_str = re.sub(r"[^\w ]", "", query) + all_tokens = nltk.word_tokenize(e_str) + stop_words = stopwords.words("english") + for token in all_tokens: + if token not in stop_words: + entities.append(token) + with self._get_cursor() as cur: + cur.execute( + f"select meta, text, embedding FROM {self.table_name}" + f" WHERE CONTAINS(text, :1, 1) > 0 order by score(1) desc fetch first {top_k} rows only", + [" ACCUM ".join(entities)], + ) + docs = [] + for record in cur: + metadata, text, embedding = record + docs.append(Document(page_content=text, vector=embedding, metadata=metadata)) + return docs + else: + return [Document(page_content="", metadata={})] + return [] + + def delete(self) -> None: + with self._get_cursor() as cur: + cur.execute(f"DROP TABLE IF EXISTS {self.table_name} cascade constraints") + + def _create_collection(self, dimension: int): + cache_key = f"vector_indexing_{self._collection_name}" + lock_name = f"{cache_key}_lock" + with redis_client.lock(lock_name, timeout=20): + collection_exist_cache_key = f"vector_indexing_{self._collection_name}" + if redis_client.get(collection_exist_cache_key): + return + + with self._get_cursor() as cur: + cur.execute(SQL_CREATE_TABLE.format(table_name=self.table_name)) + redis_client.set(collection_exist_cache_key, 1, ex=3600) + with self._get_cursor() as cur: + cur.execute(SQL_CREATE_INDEX.format(table_name=self.table_name)) + + +class OracleVectorFactory(AbstractVectorFactory): + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> OracleVector: + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id) + dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.ORACLE, collection_name)) + + return OracleVector( + collection_name=collection_name, + config=OracleVectorConfig( + host=dify_config.ORACLE_HOST or "localhost", + port=dify_config.ORACLE_PORT, + user=dify_config.ORACLE_USER or "system", + password=dify_config.ORACLE_PASSWORD or "oracle", + database=dify_config.ORACLE_DATABASE or "orcl", + ), + ) diff --git a/api/core/rag/datasource/vdb/pgvecto_rs/__init__.py b/api/core/rag/datasource/vdb/pgvecto_rs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/pgvecto_rs/collection.py b/api/core/rag/datasource/vdb/pgvecto_rs/collection.py new file mode 100644 index 0000000000000000000000000000000000000000..c335bc610d70981907604642587ddf9d242b46d2 --- /dev/null +++ b/api/core/rag/datasource/vdb/pgvecto_rs/collection.py @@ -0,0 +1,12 @@ +from uuid import UUID + +from numpy import ndarray +from sqlalchemy.orm import DeclarativeBase, Mapped + + +class CollectionORM(DeclarativeBase): + __tablename__: str + id: Mapped[UUID] + text: Mapped[str] + meta: Mapped[dict] + vector: Mapped[ndarray] diff --git a/api/core/rag/datasource/vdb/pgvecto_rs/pgvecto_rs.py b/api/core/rag/datasource/vdb/pgvecto_rs/pgvecto_rs.py new file mode 100644 index 0000000000000000000000000000000000000000..221bc68d68a6f7c2fcaa29033cb3570cd812b1bc --- /dev/null +++ b/api/core/rag/datasource/vdb/pgvecto_rs/pgvecto_rs.py @@ -0,0 +1,232 @@ +import json +import logging +from typing import Any +from uuid import UUID, uuid4 + +from numpy import ndarray +from pgvecto_rs.sqlalchemy import VECTOR # type: ignore +from pydantic import BaseModel, model_validator +from sqlalchemy import Float, String, create_engine, insert, select, text +from sqlalchemy import text as sql_text +from sqlalchemy.dialects import postgresql +from sqlalchemy.orm import Mapped, Session, mapped_column + +from configs import dify_config +from core.rag.datasource.vdb.pgvecto_rs.collection import CollectionORM +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from extensions.ext_redis import redis_client +from models.dataset import Dataset + +logger = logging.getLogger(__name__) + + +class PgvectoRSConfig(BaseModel): + host: str + port: int + user: str + password: str + database: str + + @model_validator(mode="before") + @classmethod + def validate_config(cls, values: dict) -> dict: + if not values["host"]: + raise ValueError("config PGVECTO_RS_HOST is required") + if not values["port"]: + raise ValueError("config PGVECTO_RS_PORT is required") + if not values["user"]: + raise ValueError("config PGVECTO_RS_USER is required") + if not values["password"]: + raise ValueError("config PGVECTO_RS_PASSWORD is required") + if not values["database"]: + raise ValueError("config PGVECTO_RS_DATABASE is required") + return values + + +class PGVectoRS(BaseVector): + def __init__(self, collection_name: str, config: PgvectoRSConfig, dim: int): + super().__init__(collection_name) + self._client_config = config + self._url = ( + f"postgresql+psycopg2://{config.user}:{config.password}@{config.host}:{config.port}/{config.database}" + ) + self._client = create_engine(self._url) + with Session(self._client) as session: + session.execute(text("CREATE EXTENSION IF NOT EXISTS vectors")) + session.commit() + self._fields: list[str] = [] + + class _Table(CollectionORM): + __tablename__ = collection_name + __table_args__ = {"extend_existing": True} + id: Mapped[UUID] = mapped_column( + postgresql.UUID(as_uuid=True), + primary_key=True, + ) + text: Mapped[str] = mapped_column(String) + meta: Mapped[dict] = mapped_column(postgresql.JSONB) + vector: Mapped[ndarray] = mapped_column(VECTOR(dim)) + + self._table = _Table + self._distance_op = "<=>" + + def get_type(self) -> str: + return VectorType.PGVECTO_RS + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + self.create_collection(len(embeddings[0])) + self.add_texts(texts, embeddings) + + def create_collection(self, dimension: int): + lock_name = "vector_indexing_lock_{}".format(self._collection_name) + with redis_client.lock(lock_name, timeout=20): + collection_exist_cache_key = "vector_indexing_{}".format(self._collection_name) + if redis_client.get(collection_exist_cache_key): + return + index_name = f"{self._collection_name}_embedding_index" + with Session(self._client) as session: + create_statement = sql_text(f""" + CREATE TABLE IF NOT EXISTS {self._collection_name} ( + id UUID PRIMARY KEY, + text TEXT NOT NULL, + meta JSONB NOT NULL, + vector vector({dimension}) NOT NULL + ) using heap; + """) + session.execute(create_statement) + index_statement = sql_text(f""" + CREATE INDEX IF NOT EXISTS {index_name} + ON {self._collection_name} USING vectors(vector vector_l2_ops) + WITH (options = $$ + optimizing.optimizing_threads = 30 + segment.max_growing_segment_size = 2000 + segment.max_sealed_segment_size = 30000000 + [indexing.hnsw] + m=30 + ef_construction=500 + $$); + """) + session.execute(index_statement) + session.commit() + redis_client.set(collection_exist_cache_key, 1, ex=3600) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + pks = [] + with Session(self._client) as session: + for document, embedding in zip(documents, embeddings): + pk = uuid4() + session.execute( + insert(self._table).values( + id=pk, + text=document.page_content, + meta=document.metadata, + vector=embedding, + ), + ) + pks.append(pk) + session.commit() + + return pks + + def get_ids_by_metadata_field(self, key: str, value: str): + result = None + with Session(self._client) as session: + select_statement = sql_text(f"SELECT id FROM {self._collection_name} WHERE meta->>'{key}' = '{value}'; ") + result = session.execute(select_statement).fetchall() + if result: + return [item[0] for item in result] + else: + return None + + def delete_by_metadata_field(self, key: str, value: str): + ids = self.get_ids_by_metadata_field(key, value) + if ids: + with Session(self._client) as session: + select_statement = sql_text(f"DELETE FROM {self._collection_name} WHERE id = ANY(:ids)") + session.execute(select_statement, {"ids": ids}) + session.commit() + + def delete_by_ids(self, ids: list[str]) -> None: + with Session(self._client) as session: + select_statement = sql_text( + f"SELECT id FROM {self._collection_name} WHERE meta->>'doc_id' = ANY (:doc_ids); " + ) + result = session.execute(select_statement, {"doc_ids": ids}).fetchall() + if result: + ids = [item[0] for item in result] + if ids: + with Session(self._client) as session: + select_statement = sql_text(f"DELETE FROM {self._collection_name} WHERE id = ANY(:ids)") + session.execute(select_statement, {"ids": ids}) + session.commit() + + def delete(self) -> None: + with Session(self._client) as session: + session.execute(sql_text(f"DROP TABLE IF EXISTS {self._collection_name}")) + session.commit() + + def text_exists(self, id: str) -> bool: + with Session(self._client) as session: + select_statement = sql_text( + f"SELECT id FROM {self._collection_name} WHERE meta->>'doc_id' = '{id}' limit 1; " + ) + result = session.execute(select_statement).fetchall() + return len(result) > 0 + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + with Session(self._client) as session: + stmt = ( + select( + self._table, + self._table.vector.op(self._distance_op, return_type=Float)( + query_vector, + ).label("distance"), + ) + .limit(kwargs.get("top_k", 4)) + .order_by("distance") + ) + res = session.execute(stmt) + results = [(row[0], row[1]) for row in res] + + # Organize results. + docs = [] + for record, dis in results: + metadata = record.meta + score = 1 - dis + metadata["score"] = score + score_threshold = float(kwargs.get("score_threshold") or 0.0) + if score > score_threshold: + doc = Document(page_content=record.text, metadata=metadata) + docs.append(doc) + return docs + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + return [] + + +class PGVectoRSFactory(AbstractVectorFactory): + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> PGVectoRS: + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix.lower() + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id).lower() + dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.PGVECTO_RS, collection_name)) + dim = len(embeddings.embed_query("pgvecto_rs")) + + return PGVectoRS( + collection_name=collection_name, + config=PgvectoRSConfig( + host=dify_config.PGVECTO_RS_HOST or "localhost", + port=dify_config.PGVECTO_RS_PORT or 5432, + user=dify_config.PGVECTO_RS_USER or "postgres", + password=dify_config.PGVECTO_RS_PASSWORD or "", + database=dify_config.PGVECTO_RS_DATABASE or "postgres", + ), + dim=dim, + ) diff --git a/api/core/rag/datasource/vdb/pgvector/__init__.py b/api/core/rag/datasource/vdb/pgvector/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/pgvector/pgvector.py b/api/core/rag/datasource/vdb/pgvector/pgvector.py new file mode 100644 index 0000000000000000000000000000000000000000..c8a1e4f90cd0f54905a48333da486b1a2ed60c3b --- /dev/null +++ b/api/core/rag/datasource/vdb/pgvector/pgvector.py @@ -0,0 +1,241 @@ +import json +import uuid +from contextlib import contextmanager +from typing import Any + +import psycopg2.extras # type: ignore +import psycopg2.pool # type: ignore +from pydantic import BaseModel, model_validator + +from configs import dify_config +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from extensions.ext_redis import redis_client +from models.dataset import Dataset + + +class PGVectorConfig(BaseModel): + host: str + port: int + user: str + password: str + database: str + min_connection: int + max_connection: int + + @model_validator(mode="before") + @classmethod + def validate_config(cls, values: dict) -> dict: + if not values["host"]: + raise ValueError("config PGVECTOR_HOST is required") + if not values["port"]: + raise ValueError("config PGVECTOR_PORT is required") + if not values["user"]: + raise ValueError("config PGVECTOR_USER is required") + if not values["password"]: + raise ValueError("config PGVECTOR_PASSWORD is required") + if not values["database"]: + raise ValueError("config PGVECTOR_DATABASE is required") + if not values["min_connection"]: + raise ValueError("config PGVECTOR_MIN_CONNECTION is required") + if not values["max_connection"]: + raise ValueError("config PGVECTOR_MAX_CONNECTION is required") + if values["min_connection"] > values["max_connection"]: + raise ValueError("config PGVECTOR_MIN_CONNECTION should less than PGVECTOR_MAX_CONNECTION") + return values + + +SQL_CREATE_TABLE = """ +CREATE TABLE IF NOT EXISTS {table_name} ( + id UUID PRIMARY KEY, + text TEXT NOT NULL, + meta JSONB NOT NULL, + embedding vector({dimension}) NOT NULL +) using heap; +""" + +SQL_CREATE_INDEX = """ +CREATE INDEX IF NOT EXISTS embedding_cosine_v1_idx ON {table_name} +USING hnsw (embedding vector_cosine_ops) WITH (m = 16, ef_construction = 64); +""" + + +class PGVector(BaseVector): + def __init__(self, collection_name: str, config: PGVectorConfig): + super().__init__(collection_name) + self.pool = self._create_connection_pool(config) + self.table_name = f"embedding_{collection_name}" + + def get_type(self) -> str: + return VectorType.PGVECTOR + + def _create_connection_pool(self, config: PGVectorConfig): + return psycopg2.pool.SimpleConnectionPool( + config.min_connection, + config.max_connection, + host=config.host, + port=config.port, + user=config.user, + password=config.password, + database=config.database, + ) + + @contextmanager + def _get_cursor(self): + conn = self.pool.getconn() + cur = conn.cursor() + try: + yield cur + finally: + cur.close() + conn.commit() + self.pool.putconn(conn) + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + dimension = len(embeddings[0]) + self._create_collection(dimension) + return self.add_texts(texts, embeddings) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + values = [] + pks = [] + for i, doc in enumerate(documents): + if doc.metadata is not None: + doc_id = doc.metadata.get("doc_id", str(uuid.uuid4())) + pks.append(doc_id) + values.append( + ( + doc_id, + doc.page_content, + json.dumps(doc.metadata), + embeddings[i], + ) + ) + with self._get_cursor() as cur: + psycopg2.extras.execute_values( + cur, f"INSERT INTO {self.table_name} (id, text, meta, embedding) VALUES %s", values + ) + return pks + + def text_exists(self, id: str) -> bool: + with self._get_cursor() as cur: + cur.execute(f"SELECT id FROM {self.table_name} WHERE id = %s", (id,)) + return cur.fetchone() is not None + + def get_by_ids(self, ids: list[str]) -> list[Document]: + with self._get_cursor() as cur: + cur.execute(f"SELECT meta, text FROM {self.table_name} WHERE id IN %s", (tuple(ids),)) + docs = [] + for record in cur: + docs.append(Document(page_content=record[1], metadata=record[0])) + return docs + + def delete_by_ids(self, ids: list[str]) -> None: + # Avoiding crashes caused by performing delete operations on empty lists in certain scenarios + # Scenario 1: extract a document fails, resulting in a table not being created. + # Then clicking the retry button triggers a delete operation on an empty list. + if not ids: + return + with self._get_cursor() as cur: + cur.execute(f"DELETE FROM {self.table_name} WHERE id IN %s", (tuple(ids),)) + + def delete_by_metadata_field(self, key: str, value: str) -> None: + with self._get_cursor() as cur: + cur.execute(f"DELETE FROM {self.table_name} WHERE meta->>%s = %s", (key, value)) + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + """ + Search the nearest neighbors to a vector. + + :param query_vector: The input vector to search for similar items. + :param top_k: The number of nearest neighbors to return, default is 5. + :return: List of Documents that are nearest to the query vector. + """ + top_k = kwargs.get("top_k", 4) + + with self._get_cursor() as cur: + cur.execute( + f"SELECT meta, text, embedding <=> %s AS distance FROM {self.table_name}" + f" ORDER BY distance LIMIT {top_k}", + (json.dumps(query_vector),), + ) + docs = [] + score_threshold = float(kwargs.get("score_threshold") or 0.0) + for record in cur: + metadata, text, distance = record + score = 1 - distance + metadata["score"] = score + if score > score_threshold: + docs.append(Document(page_content=text, metadata=metadata)) + return docs + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + top_k = kwargs.get("top_k", 5) + + with self._get_cursor() as cur: + cur.execute( + f"""SELECT meta, text, ts_rank(to_tsvector(coalesce(text, '')), plainto_tsquery(%s)) AS score + FROM {self.table_name} + WHERE to_tsvector(text) @@ plainto_tsquery(%s) + ORDER BY score DESC + LIMIT {top_k}""", + # f"'{query}'" is required in order to account for whitespace in query + (f"'{query}'", f"'{query}'"), + ) + + docs = [] + + for record in cur: + metadata, text, score = record + metadata["score"] = score + docs.append(Document(page_content=text, metadata=metadata)) + + return docs + + def delete(self) -> None: + with self._get_cursor() as cur: + cur.execute(f"DROP TABLE IF EXISTS {self.table_name}") + + def _create_collection(self, dimension: int): + cache_key = f"vector_indexing_{self._collection_name}" + lock_name = f"{cache_key}_lock" + with redis_client.lock(lock_name, timeout=20): + collection_exist_cache_key = f"vector_indexing_{self._collection_name}" + if redis_client.get(collection_exist_cache_key): + return + + with self._get_cursor() as cur: + cur.execute("CREATE EXTENSION IF NOT EXISTS vector") + cur.execute(SQL_CREATE_TABLE.format(table_name=self.table_name, dimension=dimension)) + # PG hnsw index only support 2000 dimension or less + # ref: https://github.com/pgvector/pgvector?tab=readme-ov-file#indexing + if dimension <= 2000: + cur.execute(SQL_CREATE_INDEX.format(table_name=self.table_name)) + redis_client.set(collection_exist_cache_key, 1, ex=3600) + + +class PGVectorFactory(AbstractVectorFactory): + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> PGVector: + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id) + dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.PGVECTOR, collection_name)) + + return PGVector( + collection_name=collection_name, + config=PGVectorConfig( + host=dify_config.PGVECTOR_HOST or "localhost", + port=dify_config.PGVECTOR_PORT, + user=dify_config.PGVECTOR_USER or "postgres", + password=dify_config.PGVECTOR_PASSWORD or "", + database=dify_config.PGVECTOR_DATABASE or "postgres", + min_connection=dify_config.PGVECTOR_MIN_CONNECTION, + max_connection=dify_config.PGVECTOR_MAX_CONNECTION, + ), + ) diff --git a/api/core/rag/datasource/vdb/qdrant/__init__.py b/api/core/rag/datasource/vdb/qdrant/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/qdrant/qdrant_vector.py b/api/core/rag/datasource/vdb/qdrant/qdrant_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..6e94cb69db309d05e596e624cd83cb278477e2e1 --- /dev/null +++ b/api/core/rag/datasource/vdb/qdrant/qdrant_vector.py @@ -0,0 +1,449 @@ +import json +import os +import uuid +from collections.abc import Generator, Iterable, Sequence +from itertools import islice +from typing import TYPE_CHECKING, Any, Optional, Union, cast + +import qdrant_client +from flask import current_app +from pydantic import BaseModel +from qdrant_client.http import models as rest +from qdrant_client.http.models import ( + FilterSelector, + HnswConfigDiff, + PayloadSchemaType, + TextIndexParams, + TextIndexType, + TokenizerType, +) +from qdrant_client.local.qdrant_local import QdrantLocal + +from configs import dify_config +from core.rag.datasource.vdb.field import Field +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from extensions.ext_database import db +from extensions.ext_redis import redis_client +from models.dataset import Dataset, DatasetCollectionBinding + +if TYPE_CHECKING: + from qdrant_client import grpc # noqa + from qdrant_client.conversions import common_types + from qdrant_client.http import models as rest + + DictFilter = dict[str, Union[str, int, bool, dict, list]] + MetadataFilter = Union[DictFilter, common_types.Filter] + + +class QdrantConfig(BaseModel): + endpoint: str + api_key: Optional[str] = None + timeout: float = 20 + root_path: Optional[str] = None + grpc_port: int = 6334 + prefer_grpc: bool = False + + def to_qdrant_params(self): + if self.endpoint and self.endpoint.startswith("path:"): + path = self.endpoint.replace("path:", "") + if not os.path.isabs(path): + if not self.root_path: + raise ValueError("Root path is not set") + path = os.path.join(self.root_path, path) + + return {"path": path} + else: + return { + "url": self.endpoint, + "api_key": self.api_key, + "timeout": self.timeout, + "verify": self.endpoint.startswith("https"), + "grpc_port": self.grpc_port, + "prefer_grpc": self.prefer_grpc, + } + + +class QdrantVector(BaseVector): + def __init__(self, collection_name: str, group_id: str, config: QdrantConfig, distance_func: str = "Cosine"): + super().__init__(collection_name) + self._client_config = config + self._client = qdrant_client.QdrantClient(**self._client_config.to_qdrant_params()) + self._distance_func = distance_func.upper() + self._group_id = group_id + + def get_type(self) -> str: + return VectorType.QDRANT + + def to_index_struct(self) -> dict: + return {"type": self.get_type(), "vector_store": {"class_prefix": self._collection_name}} + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + if texts: + # get embedding vector size + vector_size = len(embeddings[0]) + # get collection name + collection_name = self._collection_name + # create collection + self.create_collection(collection_name, vector_size) + + self.add_texts(texts, embeddings, **kwargs) + + def create_collection(self, collection_name: str, vector_size: int): + lock_name = "vector_indexing_lock_{}".format(collection_name) + with redis_client.lock(lock_name, timeout=20): + collection_exist_cache_key = "vector_indexing_{}".format(self._collection_name) + if redis_client.get(collection_exist_cache_key): + return + collection_name = collection_name or uuid.uuid4().hex + all_collection_name = [] + collections_response = self._client.get_collections() + collection_list = collections_response.collections + for collection in collection_list: + all_collection_name.append(collection.name) + if collection_name not in all_collection_name: + from qdrant_client.http import models as rest + + vectors_config = rest.VectorParams( + size=vector_size, + distance=rest.Distance[self._distance_func], + ) + hnsw_config = HnswConfigDiff( + m=0, + payload_m=16, + ef_construct=100, + full_scan_threshold=10000, + max_indexing_threads=0, + on_disk=False, + ) + self._client.recreate_collection( + collection_name=collection_name, + vectors_config=vectors_config, + hnsw_config=hnsw_config, + timeout=int(self._client_config.timeout), + ) + + # create group_id payload index + self._client.create_payload_index( + collection_name, Field.GROUP_KEY.value, field_schema=PayloadSchemaType.KEYWORD + ) + # create doc_id payload index + self._client.create_payload_index( + collection_name, Field.DOC_ID.value, field_schema=PayloadSchemaType.KEYWORD + ) + # create full text index + text_index_params = TextIndexParams( + type=TextIndexType.TEXT, + tokenizer=TokenizerType.MULTILINGUAL, + min_token_len=2, + max_token_len=20, + lowercase=True, + ) + self._client.create_payload_index( + collection_name, Field.CONTENT_KEY.value, field_schema=text_index_params + ) + redis_client.set(collection_exist_cache_key, 1, ex=3600) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + uuids = self._get_uuids(documents) + texts = [d.page_content for d in documents] + metadatas = [d.metadata for d in documents] + added_ids = [] + # Filter out None values from metadatas list to match expected type + filtered_metadatas = [m for m in metadatas if m is not None] + for batch_ids, points in self._generate_rest_batches( + texts, embeddings, filtered_metadatas, uuids, 64, self._group_id + ): + self._client.upsert(collection_name=self._collection_name, points=points) + added_ids.extend(batch_ids) + + return added_ids + + def _generate_rest_batches( + self, + texts: Iterable[str], + embeddings: list[list[float]], + metadatas: Optional[list[dict]] = None, + ids: Optional[Sequence[str]] = None, + batch_size: int = 64, + group_id: Optional[str] = None, + ) -> Generator[tuple[list[str], list[rest.PointStruct]], None, None]: + from qdrant_client.http import models as rest + + texts_iterator = iter(texts) + embeddings_iterator = iter(embeddings) + metadatas_iterator = iter(metadatas or []) + ids_iterator = iter(ids or [uuid.uuid4().hex for _ in iter(texts)]) + while batch_texts := list(islice(texts_iterator, batch_size)): + # Take the corresponding metadata and id for each text in a batch + batch_metadatas = list(islice(metadatas_iterator, batch_size)) or None + batch_ids = list(islice(ids_iterator, batch_size)) + + # Generate the embeddings for all the texts in a batch + batch_embeddings = list(islice(embeddings_iterator, batch_size)) + + points = [ + rest.PointStruct( + id=point_id, + vector=vector, + payload=payload, + ) + for point_id, vector, payload in zip( + batch_ids, + batch_embeddings, + self._build_payloads( + batch_texts, + batch_metadatas, + Field.CONTENT_KEY.value, + Field.METADATA_KEY.value, + group_id or "", # Ensure group_id is never None + Field.GROUP_KEY.value, + ), + ) + ] + + yield batch_ids, points + + @classmethod + def _build_payloads( + cls, + texts: Iterable[str], + metadatas: Optional[list[dict]], + content_payload_key: str, + metadata_payload_key: str, + group_id: str, + group_payload_key: str, + ) -> list[dict]: + payloads = [] + for i, text in enumerate(texts): + if text is None: + raise ValueError( + "At least one of the texts is None. Please remove it before " + "calling .from_texts or .add_texts on Qdrant instance." + ) + metadata = metadatas[i] if metadatas is not None else None + payloads.append({content_payload_key: text, metadata_payload_key: metadata, group_payload_key: group_id}) + + return payloads + + def delete_by_metadata_field(self, key: str, value: str): + from qdrant_client.http import models + from qdrant_client.http.exceptions import UnexpectedResponse + + try: + filter = models.Filter( + must=[ + models.FieldCondition( + key=f"metadata.{key}", + match=models.MatchValue(value=value), + ), + ], + ) + + self._reload_if_needed() + + self._client.delete( + collection_name=self._collection_name, + points_selector=FilterSelector(filter=filter), + ) + except UnexpectedResponse as e: + # Collection does not exist, so return + if e.status_code == 404: + return + # Some other error occurred, so re-raise the exception + else: + raise e + + def delete(self): + from qdrant_client.http import models + from qdrant_client.http.exceptions import UnexpectedResponse + + try: + filter = models.Filter( + must=[ + models.FieldCondition( + key="group_id", + match=models.MatchValue(value=self._group_id), + ), + ], + ) + self._client.delete( + collection_name=self._collection_name, + points_selector=FilterSelector(filter=filter), + ) + except UnexpectedResponse as e: + # Collection does not exist, so return + if e.status_code == 404: + return + # Some other error occurred, so re-raise the exception + else: + raise e + + def delete_by_ids(self, ids: list[str]) -> None: + from qdrant_client.http import models + from qdrant_client.http.exceptions import UnexpectedResponse + + for node_id in ids: + try: + filter = models.Filter( + must=[ + models.FieldCondition( + key="metadata.doc_id", + match=models.MatchValue(value=node_id), + ), + ], + ) + self._client.delete( + collection_name=self._collection_name, + points_selector=FilterSelector(filter=filter), + ) + except UnexpectedResponse as e: + # Collection does not exist, so return + if e.status_code == 404: + return + # Some other error occurred, so re-raise the exception + else: + raise e + + def text_exists(self, id: str) -> bool: + all_collection_name = [] + collections_response = self._client.get_collections() + collection_list = collections_response.collections + for collection in collection_list: + all_collection_name.append(collection.name) + if self._collection_name not in all_collection_name: + return False + response = self._client.retrieve(collection_name=self._collection_name, ids=[id]) + + return len(response) > 0 + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + from qdrant_client.http import models + + filter = models.Filter( + must=[ + models.FieldCondition( + key="group_id", + match=models.MatchValue(value=self._group_id), + ), + ], + ) + results = self._client.search( + collection_name=self._collection_name, + query_vector=query_vector, + query_filter=filter, + limit=kwargs.get("top_k", 4), + with_payload=True, + with_vectors=True, + score_threshold=float(kwargs.get("score_threshold") or 0.0), + ) + docs = [] + for result in results: + if result.payload is None: + continue + metadata = result.payload.get(Field.METADATA_KEY.value) or {} + # duplicate check score threshold + score_threshold = float(kwargs.get("score_threshold") or 0.0) + if result.score > score_threshold: + metadata["score"] = result.score + doc = Document( + page_content=result.payload.get(Field.CONTENT_KEY.value, ""), + metadata=metadata, + ) + docs.append(doc) + # Sort the documents by score in descending order + docs = sorted(docs, key=lambda x: x.metadata["score"] if x.metadata is not None else 0, reverse=True) + return docs + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + """Return docs most similar by bm25. + Returns: + List of documents most similar to the query text and distance for each. + """ + from qdrant_client.http import models + + scroll_filter = models.Filter( + must=[ + models.FieldCondition( + key="group_id", + match=models.MatchValue(value=self._group_id), + ), + models.FieldCondition( + key="page_content", + match=models.MatchText(text=query), + ), + ] + ) + response = self._client.scroll( + collection_name=self._collection_name, + scroll_filter=scroll_filter, + limit=kwargs.get("top_k", 2), + with_payload=True, + with_vectors=True, + ) + results = response[0] + documents = [] + for result in results: + if result: + document = self._document_from_scored_point(result, Field.CONTENT_KEY.value, Field.METADATA_KEY.value) + documents.append(document) + + return documents + + def _reload_if_needed(self): + if isinstance(self._client, QdrantLocal): + self._client = cast(QdrantLocal, self._client) + self._client._load() + + @classmethod + def _document_from_scored_point( + cls, + scored_point: Any, + content_payload_key: str, + metadata_payload_key: str, + ) -> Document: + return Document( + page_content=scored_point.payload.get(content_payload_key), + vector=scored_point.vector, + metadata=scored_point.payload.get(metadata_payload_key) or {}, + ) + + +class QdrantVectorFactory(AbstractVectorFactory): + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> QdrantVector: + if dataset.collection_binding_id: + dataset_collection_binding = ( + db.session.query(DatasetCollectionBinding) + .filter(DatasetCollectionBinding.id == dataset.collection_binding_id) + .one_or_none() + ) + if dataset_collection_binding: + collection_name = dataset_collection_binding.collection_name + else: + raise ValueError("Dataset Collection Bindings is not exist!") + else: + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id) + + if not dataset.index_struct_dict: + dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.QDRANT, collection_name)) + + return QdrantVector( + collection_name=collection_name, + group_id=dataset.id, + config=QdrantConfig( + endpoint=dify_config.QDRANT_URL or "", + api_key=dify_config.QDRANT_API_KEY, + root_path=str(current_app.config.root_path), + timeout=dify_config.QDRANT_CLIENT_TIMEOUT, + grpc_port=dify_config.QDRANT_GRPC_PORT, + prefer_grpc=dify_config.QDRANT_GRPC_ENABLED, + ), + ) diff --git a/api/core/rag/datasource/vdb/relyt/__init__.py b/api/core/rag/datasource/vdb/relyt/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/relyt/relyt_vector.py b/api/core/rag/datasource/vdb/relyt/relyt_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..a3a20448ff7a0a77330aa855652241b1a5c9fee7 --- /dev/null +++ b/api/core/rag/datasource/vdb/relyt/relyt_vector.py @@ -0,0 +1,313 @@ +import json +import uuid +from typing import Any, Optional + +from pydantic import BaseModel, model_validator +from sqlalchemy import Column, String, Table, create_engine, insert +from sqlalchemy import text as sql_text +from sqlalchemy.dialects.postgresql import JSON, TEXT +from sqlalchemy.orm import Session + +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from models.dataset import Dataset + +try: + from sqlalchemy.orm import declarative_base +except ImportError: + from sqlalchemy.ext.declarative import declarative_base + +from configs import dify_config +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.models.document import Document +from extensions.ext_redis import redis_client + +Base = declarative_base() # type: Any + + +class RelytConfig(BaseModel): + host: str + port: int + user: str + password: str + database: str + + @model_validator(mode="before") + @classmethod + def validate_config(cls, values: dict) -> dict: + if not values["host"]: + raise ValueError("config RELYT_HOST is required") + if not values["port"]: + raise ValueError("config RELYT_PORT is required") + if not values["user"]: + raise ValueError("config RELYT_USER is required") + if not values["password"]: + raise ValueError("config RELYT_PASSWORD is required") + if not values["database"]: + raise ValueError("config RELYT_DATABASE is required") + return values + + +class RelytVector(BaseVector): + def __init__(self, collection_name: str, config: RelytConfig, group_id: str): + super().__init__(collection_name) + self.embedding_dimension = 1536 + self._client_config = config + self._url = ( + f"postgresql+psycopg2://{config.user}:{config.password}@{config.host}:{config.port}/{config.database}" + ) + self.client = create_engine(self._url) + self._fields: list[str] = [] + self._group_id = group_id + + def get_type(self) -> str: + return VectorType.RELYT + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs) -> None: + index_params: dict[str, Any] = {} + metadatas = [d.metadata for d in texts] + self.create_collection(len(embeddings[0])) + self.embedding_dimension = len(embeddings[0]) + self.add_texts(texts, embeddings) + + def create_collection(self, dimension: int): + lock_name = "vector_indexing_lock_{}".format(self._collection_name) + with redis_client.lock(lock_name, timeout=20): + collection_exist_cache_key = "vector_indexing_{}".format(self._collection_name) + if redis_client.get(collection_exist_cache_key): + return + index_name = f"{self._collection_name}_embedding_index" + with Session(self.client) as session: + drop_statement = sql_text(f"""DROP TABLE IF EXISTS "{self._collection_name}"; """) + session.execute(drop_statement) + create_statement = sql_text(f""" + CREATE TABLE IF NOT EXISTS "{self._collection_name}" ( + id TEXT PRIMARY KEY, + document TEXT NOT NULL, + metadata JSON NOT NULL, + embedding vector({dimension}) NOT NULL + ) using heap; + """) + session.execute(create_statement) + index_statement = sql_text(f""" + CREATE INDEX {index_name} + ON "{self._collection_name}" USING vectors(embedding vector_l2_ops) + WITH (options = $$ + optimizing.optimizing_threads = 30 + segment.max_growing_segment_size = 2000 + segment.max_sealed_segment_size = 30000000 + [indexing.hnsw] + m=30 + ef_construction=500 + $$); + """) + session.execute(index_statement) + session.commit() + redis_client.set(collection_exist_cache_key, 1, ex=3600) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + from pgvecto_rs.sqlalchemy import VECTOR # type: ignore + + ids = [str(uuid.uuid1()) for _ in documents] + metadatas = [d.metadata for d in documents if d.metadata is not None] + for metadata in metadatas: + metadata["group_id"] = self._group_id + texts = [d.page_content for d in documents] + + # Define the table schema + chunks_table = Table( + self._collection_name, + Base.metadata, + Column("id", TEXT, primary_key=True), + Column("embedding", VECTOR(len(embeddings[0]))), + Column("document", String, nullable=True), + Column("metadata", JSON, nullable=True), + extend_existing=True, + ) + + chunks_table_data = [] + with self.client.connect() as conn, conn.begin(): + for document, metadata, chunk_id, embedding in zip(texts, metadatas, ids, embeddings): + chunks_table_data.append( + { + "id": chunk_id, + "embedding": embedding, + "document": document, + "metadata": metadata, + } + ) + + # Execute the batch insert when the batch size is reached + if len(chunks_table_data) == 500: + conn.execute(insert(chunks_table).values(chunks_table_data)) + # Clear the chunks_table_data list for the next batch + chunks_table_data.clear() + + # Insert any remaining records that didn't make up a full batch + if chunks_table_data: + conn.execute(insert(chunks_table).values(chunks_table_data)) + + return ids + + def get_ids_by_metadata_field(self, key: str, value: str): + result = None + with Session(self.client) as session: + select_statement = sql_text( + f"""SELECT id FROM "{self._collection_name}" WHERE metadata->>'{key}' = '{value}'; """ + ) + result = session.execute(select_statement).fetchall() + if result: + return [item[0] for item in result] + else: + return None + + def delete_by_uuids(self, ids: Optional[list[str]] = None): + """Delete by vector IDs. + + Args: + ids: List of ids to delete. + """ + from pgvecto_rs.sqlalchemy import VECTOR + + if ids is None: + raise ValueError("No ids provided to delete.") + + # Define the table schema + chunks_table = Table( + self._collection_name, + Base.metadata, + Column("id", TEXT, primary_key=True), + Column("embedding", VECTOR(self.embedding_dimension)), + Column("document", String, nullable=True), + Column("metadata", JSON, nullable=True), + extend_existing=True, + ) + + try: + with self.client.connect() as conn, conn.begin(): + delete_condition = chunks_table.c.id.in_(ids) + conn.execute(chunks_table.delete().where(delete_condition)) + return True + except Exception as e: + print("Delete operation failed:", str(e)) + return False + + def delete_by_metadata_field(self, key: str, value: str): + ids = self.get_ids_by_metadata_field(key, value) + if ids: + self.delete_by_uuids(ids) + + def delete_by_ids(self, ids: list[str]) -> None: + with Session(self.client) as session: + ids_str = ",".join(f"'{doc_id}'" for doc_id in ids) + select_statement = sql_text( + f"""SELECT id FROM "{self._collection_name}" WHERE metadata->>'doc_id' in ({ids_str}); """ + ) + result = session.execute(select_statement).fetchall() + if result: + ids = [item[0] for item in result] + self.delete_by_uuids(ids) + + def delete(self) -> None: + with Session(self.client) as session: + session.execute(sql_text(f"""DROP TABLE IF EXISTS "{self._collection_name}";""")) + session.commit() + + def text_exists(self, id: str) -> bool: + with Session(self.client) as session: + select_statement = sql_text( + f"""SELECT id FROM "{self._collection_name}" WHERE metadata->>'doc_id' = '{id}' limit 1; """ + ) + result = session.execute(select_statement).fetchall() + return len(result) > 0 + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + results = self.similarity_search_with_score_by_vector( + k=int(kwargs.get("top_k", 4)), embedding=query_vector, filter=kwargs.get("filter") + ) + + # Organize results. + docs = [] + for document, score in results: + score_threshold = float(kwargs.get("score_threshold") or 0.0) + if 1 - score > score_threshold: + docs.append(document) + return docs + + def similarity_search_with_score_by_vector( + self, + embedding: list[float], + k: int = 4, + filter: Optional[dict] = None, + ) -> list[tuple[Document, float]]: + # Add the filter if provided + + filter_condition = "" + if filter is not None: + conditions = [ + f"metadata->>{key!r} in ({', '.join(map(repr, value))})" + if len(value) > 1 + else f"metadata->>{key!r} = {value[0]!r}" + for key, value in filter.items() + ] + filter_condition = f"WHERE {' AND '.join(conditions)}" + + # Define the base query + sql_query = f""" + set vectors.enable_search_growing = on; + set vectors.enable_search_write = on; + SELECT document, metadata, embedding <-> :embedding as distance + FROM "{self._collection_name}" + {filter_condition} + ORDER BY embedding <-> :embedding + LIMIT :k + """ + + # Set up the query parameters + embedding_str = ", ".join(format(x) for x in embedding) + embedding_str = "[" + embedding_str + "]" + params = {"embedding": embedding_str, "k": k} + + # Execute the query and fetch the results + with self.client.connect() as conn: + results = conn.execute(sql_text(sql_query), params).fetchall() + + documents_with_scores = [ + ( + Document( + page_content=result.document, + metadata=result.metadata, + ), + result.distance, + ) + for result in results + ] + return documents_with_scores + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + # milvus/zilliz/relyt doesn't support bm25 search + return [] + + +class RelytVectorFactory(AbstractVectorFactory): + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> RelytVector: + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id) + dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.RELYT, collection_name)) + + return RelytVector( + collection_name=collection_name, + config=RelytConfig( + host=dify_config.RELYT_HOST or "localhost", + port=dify_config.RELYT_PORT, + user=dify_config.RELYT_USER or "", + password=dify_config.RELYT_PASSWORD or "", + database=dify_config.RELYT_DATABASE or "default", + ), + group_id=dataset.id, + ) diff --git a/api/core/rag/datasource/vdb/tencent/__init__.py b/api/core/rag/datasource/vdb/tencent/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/tencent/tencent_vector.py b/api/core/rag/datasource/vdb/tencent/tencent_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..1a4fa7b87e10869b1291e392bbdf5f18927f14f7 --- /dev/null +++ b/api/core/rag/datasource/vdb/tencent/tencent_vector.py @@ -0,0 +1,206 @@ +import json +from typing import Any, Optional + +from pydantic import BaseModel +from tcvectordb import VectorDBClient # type: ignore +from tcvectordb.model import document, enum # type: ignore +from tcvectordb.model import index as vdb_index # type: ignore +from tcvectordb.model.document import Filter # type: ignore + +from configs import dify_config +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from extensions.ext_redis import redis_client +from models.dataset import Dataset + + +class TencentConfig(BaseModel): + url: str + api_key: Optional[str] + timeout: float = 30 + username: Optional[str] + database: Optional[str] + index_type: str = "HNSW" + metric_type: str = "L2" + shard: int = 1 + replicas: int = 2 + + def to_tencent_params(self): + return {"url": self.url, "username": self.username, "key": self.api_key, "timeout": self.timeout} + + +class TencentVector(BaseVector): + field_id: str = "id" + field_vector: str = "vector" + field_text: str = "text" + field_metadata: str = "metadata" + + def __init__(self, collection_name: str, config: TencentConfig): + super().__init__(collection_name) + self._client_config = config + self._client = VectorDBClient(**self._client_config.to_tencent_params()) + self._db = self._init_database() + + def _init_database(self): + exists = False + for db in self._client.list_databases(): + if db.database_name == self._client_config.database: + exists = True + break + if exists: + return self._client.database(self._client_config.database) + else: + return self._client.create_database(database_name=self._client_config.database) + + def get_type(self) -> str: + return VectorType.TENCENT + + def to_index_struct(self) -> dict: + return {"type": self.get_type(), "vector_store": {"class_prefix": self._collection_name}} + + def _has_collection(self) -> bool: + collections = self._db.list_collections() + return any(collection.collection_name == self._collection_name for collection in collections) + + def _create_collection(self, dimension: int) -> None: + lock_name = "vector_indexing_lock_{}".format(self._collection_name) + with redis_client.lock(lock_name, timeout=20): + collection_exist_cache_key = "vector_indexing_{}".format(self._collection_name) + if redis_client.get(collection_exist_cache_key): + return + + if self._has_collection(): + return + + self.delete() + index_type = None + for k, v in enum.IndexType.__members__.items(): + if k == self._client_config.index_type: + index_type = v + if index_type is None: + raise ValueError("unsupported index_type") + metric_type = None + for k, v in enum.MetricType.__members__.items(): + if k == self._client_config.metric_type: + metric_type = v + if metric_type is None: + raise ValueError("unsupported metric_type") + params = vdb_index.HNSWParams(m=16, efconstruction=200) + index = vdb_index.Index( + vdb_index.FilterIndex(self.field_id, enum.FieldType.String, enum.IndexType.PRIMARY_KEY), + vdb_index.VectorIndex( + self.field_vector, + dimension, + index_type, + metric_type, + params, + ), + vdb_index.FilterIndex(self.field_text, enum.FieldType.String, enum.IndexType.FILTER), + vdb_index.FilterIndex(self.field_metadata, enum.FieldType.String, enum.IndexType.FILTER), + ) + + self._db.create_collection( + name=self._collection_name, + shard=self._client_config.shard, + replicas=self._client_config.replicas, + description="Collection for Dify", + index=index, + ) + redis_client.set(collection_exist_cache_key, 1, ex=3600) + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + self._create_collection(len(embeddings[0])) + self.add_texts(texts, embeddings) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + texts = [doc.page_content for doc in documents] + metadatas = [doc.metadata for doc in documents] + total_count = len(embeddings) + docs = [] + for i in range(0, total_count): + if metadatas is None: + continue + metadata = metadatas[i] or {} + doc = document.Document( + id=metadata.get("doc_id"), + vector=embeddings[i], + text=texts[i], + metadata=json.dumps(metadata), + ) + docs.append(doc) + self._db.collection(self._collection_name).upsert(docs, self._client_config.timeout) + + def text_exists(self, id: str) -> bool: + docs = self._db.collection(self._collection_name).query(document_ids=[id]) + if docs and len(docs) > 0: + return True + return False + + def delete_by_ids(self, ids: list[str]) -> None: + if not ids: + return + self._db.collection(self._collection_name).delete(document_ids=ids) + + def delete_by_metadata_field(self, key: str, value: str) -> None: + self._db.collection(self._collection_name).delete(filter=Filter(Filter.In(key, [value]))) + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + res = self._db.collection(self._collection_name).search( + vectors=[query_vector], + params=document.HNSWSearchParams(ef=kwargs.get("ef", 10)), + retrieve_vector=False, + limit=kwargs.get("top_k", 4), + timeout=self._client_config.timeout, + ) + score_threshold = float(kwargs.get("score_threshold") or 0.0) + return self._get_search_res(res, score_threshold) + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + return [] + + def _get_search_res(self, res: list | None, score_threshold: float) -> list[Document]: + docs: list[Document] = [] + if res is None or len(res) == 0: + return docs + + for result in res[0]: + meta = result.get(self.field_metadata) + if meta is not None: + meta = json.loads(meta) + score = 1 - result.get("score", 0.0) + if score > score_threshold: + meta["score"] = score + doc = Document(page_content=result.get(self.field_text), metadata=meta) + docs.append(doc) + + return docs + + def delete(self) -> None: + self._db.drop_collection(name=self._collection_name) + + +class TencentVectorFactory(AbstractVectorFactory): + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> TencentVector: + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix.lower() + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id).lower() + dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.TENCENT, collection_name)) + + return TencentVector( + collection_name=collection_name, + config=TencentConfig( + url=dify_config.TENCENT_VECTOR_DB_URL or "", + api_key=dify_config.TENCENT_VECTOR_DB_API_KEY, + timeout=dify_config.TENCENT_VECTOR_DB_TIMEOUT, + username=dify_config.TENCENT_VECTOR_DB_USERNAME, + database=dify_config.TENCENT_VECTOR_DB_DATABASE, + shard=dify_config.TENCENT_VECTOR_DB_SHARD, + replicas=dify_config.TENCENT_VECTOR_DB_REPLICAS, + ), + ) diff --git a/api/core/rag/datasource/vdb/tidb_on_qdrant/__init__.py b/api/core/rag/datasource/vdb/tidb_on_qdrant/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_entities.py b/api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_entities.py new file mode 100644 index 0000000000000000000000000000000000000000..1e62b3c58905c5c5b48c10204ef1111a1b88fc8f --- /dev/null +++ b/api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_entities.py @@ -0,0 +1,17 @@ +from typing import Optional + +from pydantic import BaseModel + + +class ClusterEntity(BaseModel): + """ + Model Config Entity. + """ + + name: str + cluster_id: str + displayName: str + region: str + spendingLimit: Optional[int] = 1000 + version: str + createdBy: str diff --git a/api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_on_qdrant_vector.py b/api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_on_qdrant_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..549f0175eb84c44f28f1911e0b2e6c752a7bd11c --- /dev/null +++ b/api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_on_qdrant_vector.py @@ -0,0 +1,530 @@ +import json +import os +import uuid +from collections.abc import Generator, Iterable, Sequence +from itertools import islice +from typing import TYPE_CHECKING, Any, Optional, Union, cast + +import qdrant_client +import requests +from flask import current_app +from pydantic import BaseModel +from qdrant_client.http import models as rest +from qdrant_client.http.models import ( + FilterSelector, + HnswConfigDiff, + PayloadSchemaType, + TextIndexParams, + TextIndexType, + TokenizerType, +) +from qdrant_client.local.qdrant_local import QdrantLocal +from requests.auth import HTTPDigestAuth + +from configs import dify_config +from core.rag.datasource.vdb.field import Field +from core.rag.datasource.vdb.tidb_on_qdrant.tidb_service import TidbService +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from extensions.ext_database import db +from extensions.ext_redis import redis_client +from models.dataset import Dataset, TidbAuthBinding + +if TYPE_CHECKING: + from qdrant_client import grpc # noqa + from qdrant_client.conversions import common_types + from qdrant_client.http import models as rest + + DictFilter = dict[str, Union[str, int, bool, dict, list]] + MetadataFilter = Union[DictFilter, common_types.Filter] + + +class TidbOnQdrantConfig(BaseModel): + endpoint: str + api_key: Optional[str] = None + timeout: float = 20 + root_path: Optional[str] = None + grpc_port: int = 6334 + prefer_grpc: bool = False + + def to_qdrant_params(self): + if self.endpoint and self.endpoint.startswith("path:"): + path = self.endpoint.replace("path:", "") + if not os.path.isabs(path): + if self.root_path: + path = os.path.join(self.root_path, path) + else: + raise ValueError("root_path is required") + + return {"path": path} + else: + return { + "url": self.endpoint, + "api_key": self.api_key, + "timeout": self.timeout, + "verify": False, + "grpc_port": self.grpc_port, + "prefer_grpc": self.prefer_grpc, + } + + +class TidbConfig(BaseModel): + api_url: str + public_key: str + private_key: str + + +class TidbOnQdrantVector(BaseVector): + def __init__(self, collection_name: str, group_id: str, config: TidbOnQdrantConfig, distance_func: str = "Cosine"): + super().__init__(collection_name) + self._client_config = config + self._client = qdrant_client.QdrantClient(**self._client_config.to_qdrant_params()) + self._distance_func = distance_func.upper() + self._group_id = group_id + + def get_type(self) -> str: + return VectorType.TIDB_ON_QDRANT + + def to_index_struct(self) -> dict: + return {"type": self.get_type(), "vector_store": {"class_prefix": self._collection_name}} + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + if texts: + # get embedding vector size + vector_size = len(embeddings[0]) + # get collection name + collection_name = self._collection_name + # create collection + self.create_collection(collection_name, vector_size) + + self.add_texts(texts, embeddings, **kwargs) + + def create_collection(self, collection_name: str, vector_size: int): + lock_name = "vector_indexing_lock_{}".format(collection_name) + with redis_client.lock(lock_name, timeout=20): + collection_exist_cache_key = "vector_indexing_{}".format(self._collection_name) + if redis_client.get(collection_exist_cache_key): + return + collection_name = collection_name or uuid.uuid4().hex + all_collection_name = [] + collections_response = self._client.get_collections() + collection_list = collections_response.collections + for collection in collection_list: + all_collection_name.append(collection.name) + if collection_name not in all_collection_name: + from qdrant_client.http import models as rest + + vectors_config = rest.VectorParams( + size=vector_size, + distance=rest.Distance[self._distance_func], + ) + hnsw_config = HnswConfigDiff( + m=0, + payload_m=16, + ef_construct=100, + full_scan_threshold=10000, + max_indexing_threads=0, + on_disk=False, + ) + self._client.recreate_collection( + collection_name=collection_name, + vectors_config=vectors_config, + hnsw_config=hnsw_config, + timeout=int(self._client_config.timeout), + ) + + # create group_id payload index + self._client.create_payload_index( + collection_name, Field.GROUP_KEY.value, field_schema=PayloadSchemaType.KEYWORD + ) + # create doc_id payload index + self._client.create_payload_index( + collection_name, Field.DOC_ID.value, field_schema=PayloadSchemaType.KEYWORD + ) + # create full text index + text_index_params = TextIndexParams( + type=TextIndexType.TEXT, + tokenizer=TokenizerType.MULTILINGUAL, + min_token_len=2, + max_token_len=20, + lowercase=True, + ) + self._client.create_payload_index( + collection_name, Field.CONTENT_KEY.value, field_schema=text_index_params + ) + redis_client.set(collection_exist_cache_key, 1, ex=3600) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + uuids = self._get_uuids(documents) + texts = [d.page_content for d in documents] + metadatas = [d.metadata for d in documents if d.metadata is not None] + + added_ids = [] + for batch_ids, points in self._generate_rest_batches(texts, embeddings, metadatas, uuids, 64, self._group_id): + self._client.upsert(collection_name=self._collection_name, points=points) + added_ids.extend(batch_ids) + + return added_ids + + def _generate_rest_batches( + self, + texts: Iterable[str], + embeddings: list[list[float]], + metadatas: Optional[list[dict]] = None, + ids: Optional[Sequence[str]] = None, + batch_size: int = 64, + group_id: Optional[str] = None, + ) -> Generator[tuple[list[str], list[rest.PointStruct]], None, None]: + from qdrant_client.http import models as rest + + texts_iterator = iter(texts) + embeddings_iterator = iter(embeddings) + metadatas_iterator = iter(metadatas or []) + ids_iterator = iter(ids or [uuid.uuid4().hex for _ in iter(texts)]) + while batch_texts := list(islice(texts_iterator, batch_size)): + # Take the corresponding metadata and id for each text in a batch + batch_metadatas = list(islice(metadatas_iterator, batch_size)) or None + batch_ids = list(islice(ids_iterator, batch_size)) + + # Generate the embeddings for all the texts in a batch + batch_embeddings = list(islice(embeddings_iterator, batch_size)) + + points = [ + rest.PointStruct( + id=point_id, + vector=vector, + payload=payload, + ) + for point_id, vector, payload in zip( + batch_ids, + batch_embeddings, + self._build_payloads( + batch_texts, + batch_metadatas, + Field.CONTENT_KEY.value, + Field.METADATA_KEY.value, + group_id or "", + Field.GROUP_KEY.value, + ), + ) + ] + + yield batch_ids, points + + @classmethod + def _build_payloads( + cls, + texts: Iterable[str], + metadatas: Optional[list[dict]], + content_payload_key: str, + metadata_payload_key: str, + group_id: str, + group_payload_key: str, + ) -> list[dict]: + payloads = [] + for i, text in enumerate(texts): + if text is None: + raise ValueError( + "At least one of the texts is None. Please remove it before " + "calling .from_texts or .add_texts on Qdrant instance." + ) + metadata = metadatas[i] if metadatas is not None else None + payloads.append({content_payload_key: text, metadata_payload_key: metadata, group_payload_key: group_id}) + + return payloads + + def delete_by_metadata_field(self, key: str, value: str): + from qdrant_client.http import models + from qdrant_client.http.exceptions import UnexpectedResponse + + try: + filter = models.Filter( + must=[ + models.FieldCondition( + key=f"metadata.{key}", + match=models.MatchValue(value=value), + ), + ], + ) + + self._reload_if_needed() + + self._client.delete( + collection_name=self._collection_name, + points_selector=FilterSelector(filter=filter), + ) + except UnexpectedResponse as e: + # Collection does not exist, so return + if e.status_code == 404: + return + # Some other error occurred, so re-raise the exception + else: + raise e + + def delete(self): + from qdrant_client.http.exceptions import UnexpectedResponse + + try: + self._client.delete_collection(collection_name=self._collection_name) + except UnexpectedResponse as e: + # Collection does not exist, so return + if e.status_code == 404: + return + # Some other error occurred, so re-raise the exception + else: + raise e + + def delete_by_ids(self, ids: list[str]) -> None: + from qdrant_client.http import models + from qdrant_client.http.exceptions import UnexpectedResponse + + for node_id in ids: + try: + filter = models.Filter( + must=[ + models.FieldCondition( + key="metadata.doc_id", + match=models.MatchValue(value=node_id), + ), + ], + ) + self._client.delete( + collection_name=self._collection_name, + points_selector=FilterSelector(filter=filter), + ) + except UnexpectedResponse as e: + # Collection does not exist, so return + if e.status_code == 404: + return + # Some other error occurred, so re-raise the exception + else: + raise e + + def text_exists(self, id: str) -> bool: + all_collection_name = [] + collections_response = self._client.get_collections() + collection_list = collections_response.collections + for collection in collection_list: + all_collection_name.append(collection.name) + if self._collection_name not in all_collection_name: + return False + response = self._client.retrieve(collection_name=self._collection_name, ids=[id]) + + return len(response) > 0 + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + from qdrant_client.http import models + + filter = models.Filter( + must=[ + models.FieldCondition( + key="group_id", + match=models.MatchValue(value=self._group_id), + ), + ], + ) + results = self._client.search( + collection_name=self._collection_name, + query_vector=query_vector, + query_filter=filter, + limit=kwargs.get("top_k", 4), + with_payload=True, + with_vectors=True, + score_threshold=kwargs.get("score_threshold", 0.0), + ) + docs = [] + for result in results: + if result.payload is None: + continue + metadata = result.payload.get(Field.METADATA_KEY.value) or {} + # duplicate check score threshold + score_threshold = kwargs.get("score_threshold") or 0.0 + if result.score > score_threshold: + metadata["score"] = result.score + doc = Document( + page_content=result.payload.get(Field.CONTENT_KEY.value, ""), + metadata=metadata, + ) + docs.append(doc) + # Sort the documents by score in descending order + docs = sorted(docs, key=lambda x: x.metadata["score"] if x.metadata is not None else 0, reverse=True) + return docs + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + """Return docs most similar by bm25. + Returns: + List of documents most similar to the query text and distance for each. + """ + from qdrant_client.http import models + + scroll_filter = models.Filter( + must=[ + models.FieldCondition( + key="page_content", + match=models.MatchText(text=query), + ) + ] + ) + response = self._client.scroll( + collection_name=self._collection_name, + scroll_filter=scroll_filter, + limit=kwargs.get("top_k", 2), + with_payload=True, + with_vectors=True, + ) + results = response[0] + documents = [] + for result in results: + if result: + document = self._document_from_scored_point(result, Field.CONTENT_KEY.value, Field.METADATA_KEY.value) + documents.append(document) + + return documents + + def _reload_if_needed(self): + if isinstance(self._client, QdrantLocal): + self._client = cast(QdrantLocal, self._client) + self._client._load() + + @classmethod + def _document_from_scored_point( + cls, + scored_point: Any, + content_payload_key: str, + metadata_payload_key: str, + ) -> Document: + return Document( + page_content=scored_point.payload.get(content_payload_key), + vector=scored_point.vector, + metadata=scored_point.payload.get(metadata_payload_key) or {}, + ) + + +class TidbOnQdrantVectorFactory(AbstractVectorFactory): + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> TidbOnQdrantVector: + tidb_auth_binding = ( + db.session.query(TidbAuthBinding).filter(TidbAuthBinding.tenant_id == dataset.tenant_id).one_or_none() + ) + if not tidb_auth_binding: + with redis_client.lock("create_tidb_serverless_cluster_lock", timeout=900): + tidb_auth_binding = ( + db.session.query(TidbAuthBinding) + .filter(TidbAuthBinding.tenant_id == dataset.tenant_id) + .one_or_none() + ) + if tidb_auth_binding: + TIDB_ON_QDRANT_API_KEY = f"{tidb_auth_binding.account}:{tidb_auth_binding.password}" + + else: + idle_tidb_auth_binding = ( + db.session.query(TidbAuthBinding) + .filter(TidbAuthBinding.active == False, TidbAuthBinding.status == "ACTIVE") + .limit(1) + .one_or_none() + ) + if idle_tidb_auth_binding: + idle_tidb_auth_binding.active = True + idle_tidb_auth_binding.tenant_id = dataset.tenant_id + db.session.commit() + TIDB_ON_QDRANT_API_KEY = f"{idle_tidb_auth_binding.account}:{idle_tidb_auth_binding.password}" + else: + new_cluster = TidbService.create_tidb_serverless_cluster( + dify_config.TIDB_PROJECT_ID or "", + dify_config.TIDB_API_URL or "", + dify_config.TIDB_IAM_API_URL or "", + dify_config.TIDB_PUBLIC_KEY or "", + dify_config.TIDB_PRIVATE_KEY or "", + dify_config.TIDB_REGION or "", + ) + new_tidb_auth_binding = TidbAuthBinding( + cluster_id=new_cluster["cluster_id"], + cluster_name=new_cluster["cluster_name"], + account=new_cluster["account"], + password=new_cluster["password"], + tenant_id=dataset.tenant_id, + active=True, + status="ACTIVE", + ) + db.session.add(new_tidb_auth_binding) + db.session.commit() + TIDB_ON_QDRANT_API_KEY = f"{new_tidb_auth_binding.account}:{new_tidb_auth_binding.password}" + else: + TIDB_ON_QDRANT_API_KEY = f"{tidb_auth_binding.account}:{tidb_auth_binding.password}" + + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id) + dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.TIDB_ON_QDRANT, collection_name)) + + config = current_app.config + + return TidbOnQdrantVector( + collection_name=collection_name, + group_id=dataset.id, + config=TidbOnQdrantConfig( + endpoint=dify_config.TIDB_ON_QDRANT_URL or "", + api_key=TIDB_ON_QDRANT_API_KEY, + root_path=str(config.root_path), + timeout=dify_config.TIDB_ON_QDRANT_CLIENT_TIMEOUT, + grpc_port=dify_config.TIDB_ON_QDRANT_GRPC_PORT, + prefer_grpc=dify_config.TIDB_ON_QDRANT_GRPC_ENABLED, + ), + ) + + def create_tidb_serverless_cluster(self, tidb_config: TidbConfig, display_name: str, region: str): + """ + Creates a new TiDB Serverless cluster. + :param tidb_config: The configuration for the TiDB Cloud API. + :param display_name: The user-friendly display name of the cluster (required). + :param region: The region where the cluster will be created (required). + + :return: The response from the API. + """ + region_object = { + "name": region, + } + + labels = { + "tidb.cloud/project": "1372813089454548012", + } + cluster_data = {"displayName": display_name, "region": region_object, "labels": labels} + + response = requests.post( + f"{tidb_config.api_url}/clusters", + json=cluster_data, + auth=HTTPDigestAuth(tidb_config.public_key, tidb_config.private_key), + ) + + if response.status_code == 200: + return response.json() + else: + response.raise_for_status() + + def change_tidb_serverless_root_password(self, tidb_config: TidbConfig, cluster_id: str, new_password: str): + """ + Changes the root password of a specific TiDB Serverless cluster. + + :param tidb_config: The configuration for the TiDB Cloud API. + :param cluster_id: The ID of the cluster for which the password is to be changed (required). + :param new_password: The new password for the root user (required). + :return: The response from the API. + """ + + body = {"password": new_password} + + response = requests.put( + f"{tidb_config.api_url}/clusters/{cluster_id}/password", + json=body, + auth=HTTPDigestAuth(tidb_config.public_key, tidb_config.private_key), + ) + + if response.status_code == 200: + return response.json() + else: + response.raise_for_status() diff --git a/api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_service.py b/api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_service.py new file mode 100644 index 0000000000000000000000000000000000000000..0a48c79511bf260c3b5fc320c7437ce7fb2189be --- /dev/null +++ b/api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_service.py @@ -0,0 +1,250 @@ +import time +import uuid + +import requests +from requests.auth import HTTPDigestAuth + +from configs import dify_config +from extensions.ext_database import db +from extensions.ext_redis import redis_client +from models.dataset import TidbAuthBinding + + +class TidbService: + @staticmethod + def create_tidb_serverless_cluster( + project_id: str, api_url: str, iam_url: str, public_key: str, private_key: str, region: str + ): + """ + Creates a new TiDB Serverless cluster. + :param project_id: The project ID of the TiDB Cloud project (required). + :param api_url: The URL of the TiDB Cloud API (required). + :param iam_url: The URL of the TiDB Cloud IAM API (required). + :param public_key: The public key for the API (required). + :param private_key: The private key for the API (required). + :param display_name: The user-friendly display name of the cluster (required). + :param region: The region where the cluster will be created (required). + + :return: The response from the API. + """ + + region_object = { + "name": region, + } + + labels = { + "tidb.cloud/project": project_id, + } + + spending_limit = { + "monthly": dify_config.TIDB_SPEND_LIMIT, + } + password = str(uuid.uuid4()).replace("-", "")[:16] + display_name = str(uuid.uuid4()).replace("-", "")[:16] + cluster_data = { + "displayName": display_name, + "region": region_object, + "labels": labels, + "spendingLimit": spending_limit, + "rootPassword": password, + } + + response = requests.post(f"{api_url}/clusters", json=cluster_data, auth=HTTPDigestAuth(public_key, private_key)) + + if response.status_code == 200: + response_data = response.json() + cluster_id = response_data["clusterId"] + retry_count = 0 + max_retries = 30 + while retry_count < max_retries: + cluster_response = TidbService.get_tidb_serverless_cluster(api_url, public_key, private_key, cluster_id) + if cluster_response["state"] == "ACTIVE": + user_prefix = cluster_response["userPrefix"] + return { + "cluster_id": cluster_id, + "cluster_name": display_name, + "account": f"{user_prefix}.root", + "password": password, + } + time.sleep(30) # wait 30 seconds before retrying + retry_count += 1 + else: + response.raise_for_status() + + @staticmethod + def delete_tidb_serverless_cluster(api_url: str, public_key: str, private_key: str, cluster_id: str): + """ + Deletes a specific TiDB Serverless cluster. + + :param api_url: The URL of the TiDB Cloud API (required). + :param public_key: The public key for the API (required). + :param private_key: The private key for the API (required). + :param cluster_id: The ID of the cluster to be deleted (required). + :return: The response from the API. + """ + + response = requests.delete(f"{api_url}/clusters/{cluster_id}", auth=HTTPDigestAuth(public_key, private_key)) + + if response.status_code == 200: + return response.json() + else: + response.raise_for_status() + + @staticmethod + def get_tidb_serverless_cluster(api_url: str, public_key: str, private_key: str, cluster_id: str): + """ + Deletes a specific TiDB Serverless cluster. + + :param api_url: The URL of the TiDB Cloud API (required). + :param public_key: The public key for the API (required). + :param private_key: The private key for the API (required). + :param cluster_id: The ID of the cluster to be deleted (required). + :return: The response from the API. + """ + + response = requests.get(f"{api_url}/clusters/{cluster_id}", auth=HTTPDigestAuth(public_key, private_key)) + + if response.status_code == 200: + return response.json() + else: + response.raise_for_status() + + @staticmethod + def change_tidb_serverless_root_password( + api_url: str, public_key: str, private_key: str, cluster_id: str, account: str, new_password: str + ): + """ + Changes the root password of a specific TiDB Serverless cluster. + + :param api_url: The URL of the TiDB Cloud API (required). + :param public_key: The public key for the API (required). + :param private_key: The private key for the API (required). + :param cluster_id: The ID of the cluster for which the password is to be changed (required).+ + :param account: The account for which the password is to be changed (required). + :param new_password: The new password for the root user (required). + :return: The response from the API. + """ + + body = {"password": new_password, "builtinRole": "role_admin", "customRoles": []} + + response = requests.patch( + f"{api_url}/clusters/{cluster_id}/sqlUsers/{account}", + json=body, + auth=HTTPDigestAuth(public_key, private_key), + ) + + if response.status_code == 200: + return response.json() + else: + response.raise_for_status() + + @staticmethod + def batch_update_tidb_serverless_cluster_status( + tidb_serverless_list: list[TidbAuthBinding], + project_id: str, + api_url: str, + iam_url: str, + public_key: str, + private_key: str, + ): + """ + Update the status of a new TiDB Serverless cluster. + :param project_id: The project ID of the TiDB Cloud project (required). + :param api_url: The URL of the TiDB Cloud API (required). + :param iam_url: The URL of the TiDB Cloud IAM API (required). + :param public_key: The public key for the API (required). + :param private_key: The private key for the API (required). + :param display_name: The user-friendly display name of the cluster (required). + :param region: The region where the cluster will be created (required). + + :return: The response from the API. + """ + tidb_serverless_list_map = {item.cluster_id: item for item in tidb_serverless_list} + cluster_ids = [item.cluster_id for item in tidb_serverless_list] + params = {"clusterIds": cluster_ids, "view": "BASIC"} + response = requests.get( + f"{api_url}/clusters:batchGet", params=params, auth=HTTPDigestAuth(public_key, private_key) + ) + + if response.status_code == 200: + response_data = response.json() + for item in response_data["clusters"]: + state = item["state"] + userPrefix = item["userPrefix"] + if state == "ACTIVE" and len(userPrefix) > 0: + cluster_info = tidb_serverless_list_map[item["clusterId"]] + cluster_info.status = "ACTIVE" + cluster_info.account = f"{userPrefix}.root" + db.session.add(cluster_info) + db.session.commit() + else: + response.raise_for_status() + + @staticmethod + def batch_create_tidb_serverless_cluster( + batch_size: int, project_id: str, api_url: str, iam_url: str, public_key: str, private_key: str, region: str + ) -> list[dict]: + """ + Creates a new TiDB Serverless cluster. + :param project_id: The project ID of the TiDB Cloud project (required). + :param api_url: The URL of the TiDB Cloud API (required). + :param iam_url: The URL of the TiDB Cloud IAM API (required). + :param public_key: The public key for the API (required). + :param private_key: The private key for the API (required). + :param display_name: The user-friendly display name of the cluster (required). + :param region: The region where the cluster will be created (required). + + :return: The response from the API. + """ + clusters = [] + for _ in range(batch_size): + region_object = { + "name": region, + } + + labels = { + "tidb.cloud/project": project_id, + } + + spending_limit = { + "monthly": dify_config.TIDB_SPEND_LIMIT, + } + password = str(uuid.uuid4()).replace("-", "")[:16] + display_name = str(uuid.uuid4()).replace("-", "") + cluster_data = { + "cluster": { + "displayName": display_name, + "region": region_object, + "labels": labels, + "spendingLimit": spending_limit, + "rootPassword": password, + } + } + cache_key = f"tidb_serverless_cluster_password:{display_name}" + redis_client.setex(cache_key, 3600, password) + clusters.append(cluster_data) + + request_body = {"requests": clusters} + response = requests.post( + f"{api_url}/clusters:batchCreate", json=request_body, auth=HTTPDigestAuth(public_key, private_key) + ) + + if response.status_code == 200: + response_data = response.json() + cluster_infos = [] + for item in response_data["clusters"]: + cache_key = f"tidb_serverless_cluster_password:{item['displayName']}" + cached_password = redis_client.get(cache_key) + if not cached_password: + continue + cluster_info = { + "cluster_id": item["clusterId"], + "cluster_name": item["displayName"], + "account": "root", + "password": cached_password.decode("utf-8"), + } + cluster_infos.append(cluster_info) + return cluster_infos + else: + response.raise_for_status() + return [] # FIXME for mypy, This line will not be reached as raise_for_status() will raise an exception diff --git a/api/core/rag/datasource/vdb/tidb_vector/__init__.py b/api/core/rag/datasource/vdb/tidb_vector/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/tidb_vector/tidb_vector.py b/api/core/rag/datasource/vdb/tidb_vector/tidb_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..be3a417390e8028aa03e3e11442eab1092723d45 --- /dev/null +++ b/api/core/rag/datasource/vdb/tidb_vector/tidb_vector.py @@ -0,0 +1,251 @@ +import json +import logging +from typing import Any + +import sqlalchemy +from pydantic import BaseModel, model_validator +from sqlalchemy import JSON, TEXT, Column, DateTime, String, Table, create_engine, insert +from sqlalchemy import text as sql_text +from sqlalchemy.orm import Session, declarative_base + +from configs import dify_config +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from extensions.ext_redis import redis_client +from models.dataset import Dataset + +logger = logging.getLogger(__name__) + + +class TiDBVectorConfig(BaseModel): + host: str + port: int + user: str + password: str + database: str + program_name: str + + @model_validator(mode="before") + @classmethod + def validate_config(cls, values: dict) -> dict: + if not values["host"]: + raise ValueError("config TIDB_VECTOR_HOST is required") + if not values["port"]: + raise ValueError("config TIDB_VECTOR_PORT is required") + if not values["user"]: + raise ValueError("config TIDB_VECTOR_USER is required") + if not values["database"]: + raise ValueError("config TIDB_VECTOR_DATABASE is required") + if not values["program_name"]: + raise ValueError("config APPLICATION_NAME is required") + return values + + +class TiDBVector(BaseVector): + def get_type(self) -> str: + return VectorType.TIDB_VECTOR + + def _table(self, dim: int) -> Table: + from tidb_vector.sqlalchemy import VectorType # type: ignore + + return Table( + self._collection_name, + self._orm_base.metadata, + Column("id", String(36), primary_key=True, nullable=False), + Column( + "vector", + VectorType(dim), + nullable=False, + comment="" if self._distance_func is None else f"hnsw(distance={self._distance_func})", + ), + Column("text", TEXT, nullable=False), + Column("meta", JSON, nullable=False), + Column("create_time", DateTime, server_default=sqlalchemy.text("CURRENT_TIMESTAMP")), + Column( + "update_time", DateTime, server_default=sqlalchemy.text("CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP") + ), + extend_existing=True, + ) + + def __init__(self, collection_name: str, config: TiDBVectorConfig, distance_func: str = "cosine"): + super().__init__(collection_name) + self._client_config = config + self._url = ( + f"mysql+pymysql://{config.user}:{config.password}@{config.host}:{config.port}/{config.database}?" + f"ssl_verify_cert=true&ssl_verify_identity=true&program_name={config.program_name}" + ) + self._distance_func = distance_func.lower() + self._engine = create_engine(self._url) + self._orm_base = declarative_base() + self._dimension = 1536 + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + logger.info("create collection and add texts, collection_name: " + self._collection_name) + self._create_collection(len(embeddings[0])) + self.add_texts(texts, embeddings) + self._dimension = len(embeddings[0]) + pass + + def _create_collection(self, dimension: int): + logger.info("_create_collection, collection_name " + self._collection_name) + lock_name = "vector_indexing_lock_{}".format(self._collection_name) + with redis_client.lock(lock_name, timeout=20): + collection_exist_cache_key = "vector_indexing_{}".format(self._collection_name) + if redis_client.get(collection_exist_cache_key): + return + with Session(self._engine) as session: + session.begin() + create_statement = sql_text(f""" + CREATE TABLE IF NOT EXISTS {self._collection_name} ( + id CHAR(36) PRIMARY KEY, + text TEXT NOT NULL, + meta JSON NOT NULL, + doc_id VARCHAR(64) AS (JSON_UNQUOTE(JSON_EXTRACT(meta, '$.doc_id'))) STORED, + KEY (doc_id), + vector VECTOR({dimension}) NOT NULL COMMENT "hnsw(distance={self._distance_func})", + create_time DATETIME DEFAULT CURRENT_TIMESTAMP, + update_time DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP + ); + """) + session.execute(create_statement) + # tidb vector not support 'CREATE/ADD INDEX' now + session.commit() + redis_client.set(collection_exist_cache_key, 1, ex=3600) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + table = self._table(len(embeddings[0])) + ids = self._get_uuids(documents) + metas = [d.metadata for d in documents] + texts = [d.page_content for d in documents] + + chunks_table_data = [] + with self._engine.connect() as conn, conn.begin(): + for id, text, meta, embedding in zip(ids, texts, metas, embeddings): + chunks_table_data.append({"id": id, "vector": embedding, "text": text, "meta": meta}) + + # Execute the batch insert when the batch size is reached + if len(chunks_table_data) == 500: + conn.execute(insert(table).values(chunks_table_data)) + # Clear the chunks_table_data list for the next batch + chunks_table_data.clear() + + # Insert any remaining records that didn't make up a full batch + if chunks_table_data: + conn.execute(insert(table).values(chunks_table_data)) + return ids + + def text_exists(self, id: str) -> bool: + result = self.get_ids_by_metadata_field("doc_id", id) + return bool(result) + + def delete_by_ids(self, ids: list[str]) -> None: + with Session(self._engine) as session: + ids_str = ",".join(f"'{doc_id}'" for doc_id in ids) + select_statement = sql_text( + f"""SELECT id FROM {self._collection_name} WHERE meta->>'$.doc_id' in ({ids_str}); """ + ) + result = session.execute(select_statement).fetchall() + if result: + ids = [item[0] for item in result] + self._delete_by_ids(ids) + + def _delete_by_ids(self, ids: list[str]) -> bool: + if ids is None: + raise ValueError("No ids provided to delete.") + table = self._table(self._dimension) + try: + with self._engine.connect() as conn, conn.begin(): + delete_condition = table.c.id.in_(ids) + conn.execute(table.delete().where(delete_condition)) + return True + except Exception as e: + print("Delete operation failed:", str(e)) + return False + + def get_ids_by_metadata_field(self, key: str, value: str): + with Session(self._engine) as session: + select_statement = sql_text( + f"""SELECT id FROM {self._collection_name} WHERE meta->>'$.{key}' = '{value}'; """ + ) + result = session.execute(select_statement).fetchall() + if result: + return [item[0] for item in result] + else: + return None + + def delete_by_metadata_field(self, key: str, value: str) -> None: + ids = self.get_ids_by_metadata_field(key, value) + if ids: + self._delete_by_ids(ids) + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + top_k = kwargs.get("top_k", 4) + score_threshold = float(kwargs.get("score_threshold") or 0.0) + filter = kwargs.get("filter") + distance = 1 - score_threshold + + query_vector_str = ", ".join(format(x) for x in query_vector) + query_vector_str = "[" + query_vector_str + "]" + logger.debug( + f"_collection_name: {self._collection_name}, score_threshold: {score_threshold}, distance: {distance}" + ) + + docs = [] + if self._distance_func == "l2": + tidb_func = "Vec_l2_distance" + elif self._distance_func == "cosine": + tidb_func = "Vec_Cosine_distance" + else: + tidb_func = "Vec_Cosine_distance" + + with Session(self._engine) as session: + select_statement = sql_text( + f"""SELECT meta, text, distance FROM ( + SELECT meta, text, {tidb_func}(vector, "{query_vector_str}") as distance + FROM {self._collection_name} + ORDER BY distance + LIMIT {top_k} + ) t WHERE distance < {distance};""" + ) + res = session.execute(select_statement) + results = [(row[0], row[1], row[2]) for row in res] + for meta, text, distance in results: + metadata = json.loads(meta) + metadata["score"] = 1 - distance + docs.append(Document(page_content=text, metadata=metadata)) + return docs + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + # tidb doesn't support bm25 search + return [] + + def delete(self) -> None: + with Session(self._engine) as session: + session.execute(sql_text(f"""DROP TABLE IF EXISTS {self._collection_name};""")) + session.commit() + + +class TiDBVectorFactory(AbstractVectorFactory): + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> TiDBVector: + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix.lower() + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id).lower() + dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.TIDB_VECTOR, collection_name)) + + return TiDBVector( + collection_name=collection_name, + config=TiDBVectorConfig( + host=dify_config.TIDB_VECTOR_HOST or "", + port=dify_config.TIDB_VECTOR_PORT or 0, + user=dify_config.TIDB_VECTOR_USER or "", + password=dify_config.TIDB_VECTOR_PASSWORD or "", + database=dify_config.TIDB_VECTOR_DATABASE or "", + program_name=dify_config.APPLICATION_NAME, + ), + ) diff --git a/api/core/rag/datasource/vdb/upstash/__init__.py b/api/core/rag/datasource/vdb/upstash/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/upstash/upstash_vector.py b/api/core/rag/datasource/vdb/upstash/upstash_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..5c3fee98a9633b994a96c1b74b0dd795bdf6f7a7 --- /dev/null +++ b/api/core/rag/datasource/vdb/upstash/upstash_vector.py @@ -0,0 +1,130 @@ +import json +from typing import Any +from uuid import uuid4 + +from pydantic import BaseModel, model_validator +from upstash_vector import Index, Vector + +from configs import dify_config +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from models.dataset import Dataset + + +class UpstashVectorConfig(BaseModel): + url: str + token: str + + @model_validator(mode="before") + @classmethod + def validate_config(cls, values: dict) -> dict: + if not values["url"]: + raise ValueError("Upstash URL is required") + if not values["token"]: + raise ValueError("Upstash Token is required") + return values + + +class UpstashVector(BaseVector): + def __init__(self, collection_name: str, config: UpstashVectorConfig): + super().__init__(collection_name) + self._table_name = collection_name + self.index = Index(url=config.url, token=config.token) + + def _get_index_dimension(self) -> int: + index_info = self.index.info() + if index_info and index_info.dimension: + return index_info.dimension + else: + return 1536 + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + self.add_texts(texts, embeddings) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + vectors = [ + Vector( + id=str(uuid4()), + vector=embedding, + metadata=doc.metadata, + data=doc.page_content, + ) + for doc, embedding in zip(documents, embeddings) + ] + self.index.upsert(vectors=vectors) + + def text_exists(self, id: str) -> bool: + response = self.get_ids_by_metadata_field("doc_id", id) + return len(response) > 0 + + def delete_by_ids(self, ids: list[str]) -> None: + item_ids = [] + for doc_id in ids: + ids = self.get_ids_by_metadata_field("doc_id", doc_id) + if ids: + item_ids += ids + self._delete_by_ids(ids=item_ids) + + def _delete_by_ids(self, ids: list[str]) -> None: + if ids: + self.index.delete(ids=ids) + + def get_ids_by_metadata_field(self, key: str, value: str) -> list[str]: + query_result = self.index.query( + vector=[1.001 * i for i in range(self._get_index_dimension())], + include_metadata=True, + top_k=1000, + filter=f"{key} = '{value}'", + ) + return [result.id for result in query_result] + + def delete_by_metadata_field(self, key: str, value: str) -> None: + ids = self.get_ids_by_metadata_field(key, value) + if ids: + self._delete_by_ids(ids) + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + top_k = kwargs.get("top_k", 4) + result = self.index.query(vector=query_vector, top_k=top_k, include_metadata=True, include_data=True) + docs = [] + score_threshold = float(kwargs.get("score_threshold") or 0.0) + for record in result: + metadata = record.metadata + text = record.data + score = record.score + if metadata is not None and text is not None: + metadata["score"] = score + if score > score_threshold: + docs.append(Document(page_content=text, metadata=metadata)) + return docs + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + return [] + + def delete(self) -> None: + self.index.reset() + + def get_type(self) -> str: + return VectorType.UPSTASH + + +class UpstashVectorFactory(AbstractVectorFactory): + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> UpstashVector: + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix.lower() + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id).lower() + dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.UPSTASH, collection_name)) + + return UpstashVector( + collection_name=collection_name, + config=UpstashVectorConfig( + url=dify_config.UPSTASH_VECTOR_URL or "", + token=dify_config.UPSTASH_VECTOR_TOKEN or "", + ), + ) diff --git a/api/core/rag/datasource/vdb/vector_base.py b/api/core/rag/datasource/vdb/vector_base.py new file mode 100644 index 0000000000000000000000000000000000000000..edfce2edd896eee87bc11e9cbbb50a2c04f44078 --- /dev/null +++ b/api/core/rag/datasource/vdb/vector_base.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + +from core.rag.models.document import Document + + +class BaseVector(ABC): + def __init__(self, collection_name: str): + self._collection_name = collection_name + + @abstractmethod + def get_type(self) -> str: + raise NotImplementedError + + @abstractmethod + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + raise NotImplementedError + + @abstractmethod + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + raise NotImplementedError + + @abstractmethod + def text_exists(self, id: str) -> bool: + raise NotImplementedError + + @abstractmethod + def delete_by_ids(self, ids: list[str]) -> None: + raise NotImplementedError + + def get_ids_by_metadata_field(self, key: str, value: str): + raise NotImplementedError + + @abstractmethod + def delete_by_metadata_field(self, key: str, value: str) -> None: + raise NotImplementedError + + @abstractmethod + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + raise NotImplementedError + + @abstractmethod + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + raise NotImplementedError + + @abstractmethod + def delete(self) -> None: + raise NotImplementedError + + def _filter_duplicate_texts(self, texts: list[Document]) -> list[Document]: + for text in texts.copy(): + if text.metadata and "doc_id" in text.metadata: + doc_id = text.metadata["doc_id"] + exists_duplicate_node = self.text_exists(doc_id) + if exists_duplicate_node: + texts.remove(text) + + return texts + + def _get_uuids(self, texts: list[Document]) -> list[str]: + return [text.metadata["doc_id"] for text in texts if text.metadata and "doc_id" in text.metadata] + + @property + def collection_name(self): + return self._collection_name diff --git a/api/core/rag/datasource/vdb/vector_factory.py b/api/core/rag/datasource/vdb/vector_factory.py new file mode 100644 index 0000000000000000000000000000000000000000..bdc40e29c7e83884c77e7fbc01586ff16ebb5a1c --- /dev/null +++ b/api/core/rag/datasource/vdb/vector_factory.py @@ -0,0 +1,218 @@ +from abc import ABC, abstractmethod +from typing import Any, Optional + +from configs import dify_config +from core.model_manager import ModelManager +from core.model_runtime.entities.model_entities import ModelType +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.cached_embedding import CacheEmbedding +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from extensions.ext_database import db +from extensions.ext_redis import redis_client +from models.dataset import Dataset, Whitelist + + +class AbstractVectorFactory(ABC): + @abstractmethod + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> BaseVector: + raise NotImplementedError + + @staticmethod + def gen_index_struct_dict(vector_type: VectorType, collection_name: str) -> dict: + index_struct_dict = {"type": vector_type, "vector_store": {"class_prefix": collection_name}} + return index_struct_dict + + +class Vector: + def __init__(self, dataset: Dataset, attributes: Optional[list] = None): + if attributes is None: + attributes = ["doc_id", "dataset_id", "document_id", "doc_hash"] + self._dataset = dataset + self._embeddings = self._get_embeddings() + self._attributes = attributes + self._vector_processor = self._init_vector() + + def _init_vector(self) -> BaseVector: + vector_type = dify_config.VECTOR_STORE + + if self._dataset.index_struct_dict: + vector_type = self._dataset.index_struct_dict["type"] + else: + if dify_config.VECTOR_STORE_WHITELIST_ENABLE: + whitelist = ( + db.session.query(Whitelist) + .filter(Whitelist.tenant_id == self._dataset.tenant_id, Whitelist.category == "vector_db") + .one_or_none() + ) + if whitelist: + vector_type = VectorType.TIDB_ON_QDRANT + + if not vector_type: + raise ValueError("Vector store must be specified.") + + vector_factory_cls = self.get_vector_factory(vector_type) + return vector_factory_cls().init_vector(self._dataset, self._attributes, self._embeddings) + + @staticmethod + def get_vector_factory(vector_type: str) -> type[AbstractVectorFactory]: + match vector_type: + case VectorType.CHROMA: + from core.rag.datasource.vdb.chroma.chroma_vector import ChromaVectorFactory + + return ChromaVectorFactory + case VectorType.MILVUS: + from core.rag.datasource.vdb.milvus.milvus_vector import MilvusVectorFactory + + return MilvusVectorFactory + case VectorType.MYSCALE: + from core.rag.datasource.vdb.myscale.myscale_vector import MyScaleVectorFactory + + return MyScaleVectorFactory + case VectorType.PGVECTOR: + from core.rag.datasource.vdb.pgvector.pgvector import PGVectorFactory + + return PGVectorFactory + case VectorType.PGVECTO_RS: + from core.rag.datasource.vdb.pgvecto_rs.pgvecto_rs import PGVectoRSFactory + + return PGVectoRSFactory + case VectorType.QDRANT: + from core.rag.datasource.vdb.qdrant.qdrant_vector import QdrantVectorFactory + + return QdrantVectorFactory + case VectorType.RELYT: + from core.rag.datasource.vdb.relyt.relyt_vector import RelytVectorFactory + + return RelytVectorFactory + case VectorType.ELASTICSEARCH: + from core.rag.datasource.vdb.elasticsearch.elasticsearch_vector import ElasticSearchVectorFactory + + return ElasticSearchVectorFactory + case VectorType.ELASTICSEARCH_JA: + from core.rag.datasource.vdb.elasticsearch.elasticsearch_ja_vector import ( + ElasticSearchJaVectorFactory, + ) + + return ElasticSearchJaVectorFactory + case VectorType.TIDB_VECTOR: + from core.rag.datasource.vdb.tidb_vector.tidb_vector import TiDBVectorFactory + + return TiDBVectorFactory + case VectorType.WEAVIATE: + from core.rag.datasource.vdb.weaviate.weaviate_vector import WeaviateVectorFactory + + return WeaviateVectorFactory + case VectorType.TENCENT: + from core.rag.datasource.vdb.tencent.tencent_vector import TencentVectorFactory + + return TencentVectorFactory + case VectorType.ORACLE: + from core.rag.datasource.vdb.oracle.oraclevector import OracleVectorFactory + + return OracleVectorFactory + case VectorType.OPENSEARCH: + from core.rag.datasource.vdb.opensearch.opensearch_vector import OpenSearchVectorFactory + + return OpenSearchVectorFactory + case VectorType.ANALYTICDB: + from core.rag.datasource.vdb.analyticdb.analyticdb_vector import AnalyticdbVectorFactory + + return AnalyticdbVectorFactory + case VectorType.COUCHBASE: + from core.rag.datasource.vdb.couchbase.couchbase_vector import CouchbaseVectorFactory + + return CouchbaseVectorFactory + case VectorType.BAIDU: + from core.rag.datasource.vdb.baidu.baidu_vector import BaiduVectorFactory + + return BaiduVectorFactory + case VectorType.VIKINGDB: + from core.rag.datasource.vdb.vikingdb.vikingdb_vector import VikingDBVectorFactory + + return VikingDBVectorFactory + case VectorType.UPSTASH: + from core.rag.datasource.vdb.upstash.upstash_vector import UpstashVectorFactory + + return UpstashVectorFactory + case VectorType.TIDB_ON_QDRANT: + from core.rag.datasource.vdb.tidb_on_qdrant.tidb_on_qdrant_vector import TidbOnQdrantVectorFactory + + return TidbOnQdrantVectorFactory + case VectorType.LINDORM: + from core.rag.datasource.vdb.lindorm.lindorm_vector import LindormVectorStoreFactory + + return LindormVectorStoreFactory + case VectorType.OCEANBASE: + from core.rag.datasource.vdb.oceanbase.oceanbase_vector import OceanBaseVectorFactory + + return OceanBaseVectorFactory + case _: + raise ValueError(f"Vector store {vector_type} is not supported.") + + def create(self, texts: Optional[list] = None, **kwargs): + if texts: + embeddings = self._embeddings.embed_documents([document.page_content for document in texts]) + self._vector_processor.create(texts=texts, embeddings=embeddings, **kwargs) + + def add_texts(self, documents: list[Document], **kwargs): + if kwargs.get("duplicate_check", False): + documents = self._filter_duplicate_texts(documents) + + embeddings = self._embeddings.embed_documents([document.page_content for document in documents]) + self._vector_processor.create(texts=documents, embeddings=embeddings, **kwargs) + + def text_exists(self, id: str) -> bool: + return self._vector_processor.text_exists(id) + + def delete_by_ids(self, ids: list[str]) -> None: + self._vector_processor.delete_by_ids(ids) + + def delete_by_metadata_field(self, key: str, value: str) -> None: + self._vector_processor.delete_by_metadata_field(key, value) + + def search_by_vector(self, query: str, **kwargs: Any) -> list[Document]: + query_vector = self._embeddings.embed_query(query) + return self._vector_processor.search_by_vector(query_vector, **kwargs) + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + return self._vector_processor.search_by_full_text(query, **kwargs) + + def delete(self) -> None: + self._vector_processor.delete() + # delete collection redis cache + if self._vector_processor.collection_name: + collection_exist_cache_key = "vector_indexing_{}".format(self._vector_processor.collection_name) + redis_client.delete(collection_exist_cache_key) + + def _get_embeddings(self) -> Embeddings: + model_manager = ModelManager() + + embedding_model = model_manager.get_model_instance( + tenant_id=self._dataset.tenant_id, + provider=self._dataset.embedding_model_provider, + model_type=ModelType.TEXT_EMBEDDING, + model=self._dataset.embedding_model, + ) + return CacheEmbedding(embedding_model) + + def _filter_duplicate_texts(self, texts: list[Document]) -> list[Document]: + for text in texts.copy(): + if text.metadata is None: + continue + doc_id = text.metadata["doc_id"] + if doc_id: + exists_duplicate_node = self.text_exists(doc_id) + if exists_duplicate_node: + texts.remove(text) + + return texts + + def __getattr__(self, name): + if self._vector_processor is not None: + method = getattr(self._vector_processor, name) + if callable(method): + return method + + raise AttributeError(f"'vector_processor' object has no attribute '{name}'") diff --git a/api/core/rag/datasource/vdb/vector_type.py b/api/core/rag/datasource/vdb/vector_type.py new file mode 100644 index 0000000000000000000000000000000000000000..e73411aa0d38a943a9055fefd61fe8c008beba49 --- /dev/null +++ b/api/core/rag/datasource/vdb/vector_type.py @@ -0,0 +1,26 @@ +from enum import StrEnum + + +class VectorType(StrEnum): + ANALYTICDB = "analyticdb" + CHROMA = "chroma" + MILVUS = "milvus" + MYSCALE = "myscale" + PGVECTOR = "pgvector" + PGVECTO_RS = "pgvecto-rs" + QDRANT = "qdrant" + RELYT = "relyt" + TIDB_VECTOR = "tidb_vector" + WEAVIATE = "weaviate" + OPENSEARCH = "opensearch" + TENCENT = "tencent" + ORACLE = "oracle" + ELASTICSEARCH = "elasticsearch" + ELASTICSEARCH_JA = "elasticsearch-ja" + LINDORM = "lindorm" + COUCHBASE = "couchbase" + BAIDU = "baidu" + VIKINGDB = "vikingdb" + UPSTASH = "upstash" + TIDB_ON_QDRANT = "tidb_on_qdrant" + OCEANBASE = "oceanbase" diff --git a/api/core/rag/datasource/vdb/vikingdb/__init__.py b/api/core/rag/datasource/vdb/vikingdb/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/vikingdb/vikingdb_vector.py b/api/core/rag/datasource/vdb/vikingdb/vikingdb_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..9de8761a91ca6857590cd342c5e347681083a82e --- /dev/null +++ b/api/core/rag/datasource/vdb/vikingdb/vikingdb_vector.py @@ -0,0 +1,240 @@ +import json +from typing import Any + +from pydantic import BaseModel +from volcengine.viking_db import ( # type: ignore + Data, + DistanceType, + Field, + FieldType, + IndexType, + QuantType, + VectorIndexParams, + VikingDBService, +) + +from configs import dify_config +from core.rag.datasource.vdb.field import Field as vdb_Field +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from extensions.ext_redis import redis_client +from models.dataset import Dataset + + +class VikingDBConfig(BaseModel): + access_key: str + secret_key: str + host: str + region: str + scheme: str + connection_timeout: int + socket_timeout: int + index_type: str = IndexType.HNSW + distance: str = DistanceType.L2 + quant: str = QuantType.Float + + +class VikingDBVector(BaseVector): + def __init__(self, collection_name: str, group_id: str, config: VikingDBConfig): + super().__init__(collection_name) + self._group_id = group_id + self._client_config = config + self._index_name = f"{self._collection_name}_idx" + self._client = VikingDBService( + host=config.host, + region=config.region, + scheme=config.scheme, + connection_timeout=config.connection_timeout, + socket_timeout=config.socket_timeout, + ak=config.access_key, + sk=config.secret_key, + ) + + def _has_collection(self) -> bool: + try: + self._client.get_collection(self._collection_name) + except Exception: + return False + return True + + def _has_index(self) -> bool: + try: + self._client.get_index(self._collection_name, self._index_name) + except Exception: + return False + return True + + def _create_collection(self, dimension: int): + lock_name = f"vector_indexing_lock_{self._collection_name}" + with redis_client.lock(lock_name, timeout=20): + collection_exist_cache_key = f"vector_indexing_{self._collection_name}" + if redis_client.get(collection_exist_cache_key): + return + + if not self._has_collection(): + fields = [ + Field(field_name=vdb_Field.PRIMARY_KEY.value, field_type=FieldType.String, is_primary_key=True), + Field(field_name=vdb_Field.METADATA_KEY.value, field_type=FieldType.String), + Field(field_name=vdb_Field.GROUP_KEY.value, field_type=FieldType.String), + Field(field_name=vdb_Field.CONTENT_KEY.value, field_type=FieldType.Text), + Field(field_name=vdb_Field.VECTOR.value, field_type=FieldType.Vector, dim=dimension), + ] + + self._client.create_collection( + collection_name=self._collection_name, + fields=fields, + description="Collection For Dify", + ) + + if not self._has_index(): + vector_index = VectorIndexParams( + distance=self._client_config.distance, + index_type=self._client_config.index_type, + quant=self._client_config.quant, + ) + + self._client.create_index( + collection_name=self._collection_name, + index_name=self._index_name, + vector_index=vector_index, + partition_by=vdb_Field.GROUP_KEY.value, + description="Index For Dify", + ) + redis_client.set(collection_exist_cache_key, 1, ex=3600) + + def get_type(self) -> str: + return VectorType.VIKINGDB + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + dimension = len(embeddings[0]) + self._create_collection(dimension) + self.add_texts(texts, embeddings, **kwargs) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + page_contents = [doc.page_content for doc in documents] + metadatas = [doc.metadata for doc in documents] + docs = [] + + for i, page_content in enumerate(page_contents): + metadata = {} + if metadatas is not None: + for key, val in (metadatas[i] or {}).items(): + metadata[key] = val + # FIXME: fix the type of metadata later + doc = Data( + { + vdb_Field.PRIMARY_KEY.value: metadatas[i]["doc_id"], # type: ignore + vdb_Field.VECTOR.value: embeddings[i] if embeddings else None, + vdb_Field.CONTENT_KEY.value: page_content, + vdb_Field.METADATA_KEY.value: json.dumps(metadata), + vdb_Field.GROUP_KEY.value: self._group_id, + } + ) + docs.append(doc) + + self._client.get_collection(self._collection_name).upsert_data(docs) + + def text_exists(self, id: str) -> bool: + docs = self._client.get_collection(self._collection_name).fetch_data(id) + not_exists_str = "data does not exist" + if docs is not None and not_exists_str not in docs.fields.get("message", ""): + return True + return False + + def delete_by_ids(self, ids: list[str]) -> None: + self._client.get_collection(self._collection_name).delete_data(ids) + + def get_ids_by_metadata_field(self, key: str, value: str): + # Note: Metadata field value is an dict, but vikingdb field + # not support json type + results = self._client.get_index(self._collection_name, self._index_name).search( + filter={"op": "must", "field": vdb_Field.GROUP_KEY.value, "conds": [self._group_id]}, + # max value is 5000 + limit=5000, + ) + + if not results: + return [] + + ids = [] + for result in results: + metadata = result.fields.get(vdb_Field.METADATA_KEY.value) + if metadata is not None: + metadata = json.loads(metadata) + if metadata.get(key) == value: + ids.append(result.id) + return ids + + def delete_by_metadata_field(self, key: str, value: str) -> None: + ids = self.get_ids_by_metadata_field(key, value) + self.delete_by_ids(ids) + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + results = self._client.get_index(self._collection_name, self._index_name).search_by_vector( + query_vector, limit=kwargs.get("top_k", 4) + ) + score_threshold = float(kwargs.get("score_threshold") or 0.0) + return self._get_search_res(results, score_threshold) + + def _get_search_res(self, results, score_threshold) -> list[Document]: + if len(results) == 0: + return [] + + docs = [] + for result in results: + metadata = result.fields.get(vdb_Field.METADATA_KEY.value) + if metadata is not None: + metadata = json.loads(metadata) + if result.score > score_threshold: + metadata["score"] = result.score + doc = Document(page_content=result.fields.get(vdb_Field.CONTENT_KEY.value), metadata=metadata) + docs.append(doc) + docs = sorted(docs, key=lambda x: x.metadata.get("score", 0) if x.metadata else 0, reverse=True) + return docs + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + return [] + + def delete(self) -> None: + if self._has_index(): + self._client.drop_index(self._collection_name, self._index_name) + if self._has_collection(): + self._client.drop_collection(self._collection_name) + + +class VikingDBVectorFactory(AbstractVectorFactory): + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> VikingDBVector: + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix.lower() + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id).lower() + dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.VIKINGDB, collection_name)) + + if dify_config.VIKINGDB_ACCESS_KEY is None: + raise ValueError("VIKINGDB_ACCESS_KEY should not be None") + if dify_config.VIKINGDB_SECRET_KEY is None: + raise ValueError("VIKINGDB_SECRET_KEY should not be None") + if dify_config.VIKINGDB_HOST is None: + raise ValueError("VIKINGDB_HOST should not be None") + if dify_config.VIKINGDB_REGION is None: + raise ValueError("VIKINGDB_REGION should not be None") + if dify_config.VIKINGDB_SCHEME is None: + raise ValueError("VIKINGDB_SCHEME should not be None") + return VikingDBVector( + collection_name=collection_name, + group_id=dataset.id, + config=VikingDBConfig( + access_key=dify_config.VIKINGDB_ACCESS_KEY, + secret_key=dify_config.VIKINGDB_SECRET_KEY, + host=dify_config.VIKINGDB_HOST, + region=dify_config.VIKINGDB_REGION, + scheme=dify_config.VIKINGDB_SCHEME, + connection_timeout=dify_config.VIKINGDB_CONNECTION_TIMEOUT, + socket_timeout=dify_config.VIKINGDB_SOCKET_TIMEOUT, + ), + ) diff --git a/api/core/rag/datasource/vdb/weaviate/__init__.py b/api/core/rag/datasource/vdb/weaviate/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/datasource/vdb/weaviate/weaviate_vector.py b/api/core/rag/datasource/vdb/weaviate/weaviate_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..68d043a19f171f7e863eebb595487ebbfcaf0e76 --- /dev/null +++ b/api/core/rag/datasource/vdb/weaviate/weaviate_vector.py @@ -0,0 +1,285 @@ +import datetime +import json +from typing import Any, Optional + +import requests +import weaviate # type: ignore +from pydantic import BaseModel, model_validator + +from configs import dify_config +from core.rag.datasource.vdb.field import Field +from core.rag.datasource.vdb.vector_base import BaseVector +from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory +from core.rag.datasource.vdb.vector_type import VectorType +from core.rag.embedding.embedding_base import Embeddings +from core.rag.models.document import Document +from extensions.ext_redis import redis_client +from models.dataset import Dataset + + +class WeaviateConfig(BaseModel): + endpoint: str + api_key: Optional[str] = None + batch_size: int = 100 + + @model_validator(mode="before") + @classmethod + def validate_config(cls, values: dict) -> dict: + if not values["endpoint"]: + raise ValueError("config WEAVIATE_ENDPOINT is required") + return values + + +class WeaviateVector(BaseVector): + def __init__(self, collection_name: str, config: WeaviateConfig, attributes: list): + super().__init__(collection_name) + self._client = self._init_client(config) + self._attributes = attributes + + def _init_client(self, config: WeaviateConfig) -> weaviate.Client: + auth_config = weaviate.auth.AuthApiKey(api_key=config.api_key) + + weaviate.connect.connection.has_grpc = False + + try: + client = weaviate.Client( + url=config.endpoint, auth_client_secret=auth_config, timeout_config=(5, 60), startup_period=None + ) + except requests.exceptions.ConnectionError: + raise ConnectionError("Vector database connection error") + + client.batch.configure( + # `batch_size` takes an `int` value to enable auto-batching + # (`None` is used for manual batching) + batch_size=config.batch_size, + # dynamically update the `batch_size` based on import speed + dynamic=True, + # `timeout_retries` takes an `int` value to retry on time outs + timeout_retries=3, + ) + + return client + + def get_type(self) -> str: + return VectorType.WEAVIATE + + def get_collection_name(self, dataset: Dataset) -> str: + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + if not class_prefix.endswith("_Node"): + # original class_prefix + class_prefix += "_Node" + + return class_prefix + + dataset_id = dataset.id + return Dataset.gen_collection_name_by_id(dataset_id) + + def to_index_struct(self) -> dict: + return {"type": self.get_type(), "vector_store": {"class_prefix": self._collection_name}} + + def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): + # create collection + self._create_collection() + # create vector + self.add_texts(texts, embeddings) + + def _create_collection(self): + lock_name = "vector_indexing_lock_{}".format(self._collection_name) + with redis_client.lock(lock_name, timeout=20): + collection_exist_cache_key = "vector_indexing_{}".format(self._collection_name) + if redis_client.get(collection_exist_cache_key): + return + schema = self._default_schema(self._collection_name) + if not self._client.schema.contains(schema): + # create collection + self._client.schema.create_class(schema) + redis_client.set(collection_exist_cache_key, 1, ex=3600) + + def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): + uuids = self._get_uuids(documents) + texts = [d.page_content for d in documents] + metadatas = [d.metadata for d in documents] + + ids = [] + + with self._client.batch as batch: + for i, text in enumerate(texts): + data_properties = {Field.TEXT_KEY.value: text} + if metadatas is not None: + # metadata maybe None + for key, val in (metadatas[i] or {}).items(): + data_properties[key] = self._json_serializable(val) + + batch.add_data_object( + data_object=data_properties, + class_name=self._collection_name, + uuid=uuids[i], + vector=embeddings[i] if embeddings else None, + ) + ids.append(uuids[i]) + return ids + + def delete_by_metadata_field(self, key: str, value: str): + # check whether the index already exists + schema = self._default_schema(self._collection_name) + if self._client.schema.contains(schema): + where_filter = {"operator": "Equal", "path": [key], "valueText": value} + + self._client.batch.delete_objects(class_name=self._collection_name, where=where_filter, output="minimal") + + def delete(self): + # check whether the index already exists + schema = self._default_schema(self._collection_name) + if self._client.schema.contains(schema): + self._client.schema.delete_class(self._collection_name) + + def text_exists(self, id: str) -> bool: + collection_name = self._collection_name + schema = self._default_schema(self._collection_name) + + # check whether the index already exists + if not self._client.schema.contains(schema): + return False + result = ( + self._client.query.get(collection_name) + .with_additional(["id"]) + .with_where( + { + "path": ["doc_id"], + "operator": "Equal", + "valueText": id, + } + ) + .with_limit(1) + .do() + ) + + if "errors" in result: + raise ValueError(f"Error during query: {result['errors']}") + + entries = result["data"]["Get"][collection_name] + if len(entries) == 0: + return False + + return True + + def delete_by_ids(self, ids: list[str]) -> None: + # check whether the index already exists + schema = self._default_schema(self._collection_name) + if self._client.schema.contains(schema): + for uuid in ids: + try: + self._client.data_object.delete( + class_name=self._collection_name, + uuid=uuid, + ) + except weaviate.UnexpectedStatusCodeException as e: + # tolerate not found error + if e.status_code != 404: + raise e + + def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: + """Look up similar documents by embedding vector in Weaviate.""" + collection_name = self._collection_name + properties = self._attributes + properties.append(Field.TEXT_KEY.value) + query_obj = self._client.query.get(collection_name, properties) + + vector = {"vector": query_vector} + if kwargs.get("where_filter"): + query_obj = query_obj.with_where(kwargs.get("where_filter")) + result = ( + query_obj.with_near_vector(vector) + .with_limit(kwargs.get("top_k", 4)) + .with_additional(["vector", "distance"]) + .do() + ) + if "errors" in result: + raise ValueError(f"Error during query: {result['errors']}") + + docs_and_scores = [] + for res in result["data"]["Get"][collection_name]: + text = res.pop(Field.TEXT_KEY.value) + score = 1 - res["_additional"]["distance"] + docs_and_scores.append((Document(page_content=text, metadata=res), score)) + + docs = [] + for doc, score in docs_and_scores: + score_threshold = float(kwargs.get("score_threshold") or 0.0) + # check score threshold + if score > score_threshold: + if doc.metadata is not None: + doc.metadata["score"] = score + docs.append(doc) + # Sort the documents by score in descending order + docs = sorted(docs, key=lambda x: x.metadata.get("score", 0) if x.metadata else 0, reverse=True) + return docs + + def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: + """Return docs using BM25F. + + Args: + query: Text to look up documents similar to. + k: Number of Documents to return. Defaults to 4. + + Returns: + List of Documents most similar to the query. + """ + collection_name = self._collection_name + content: dict[str, Any] = {"concepts": [query]} + properties = self._attributes + properties.append(Field.TEXT_KEY.value) + if kwargs.get("search_distance"): + content["certainty"] = kwargs.get("search_distance") + query_obj = self._client.query.get(collection_name, properties) + if kwargs.get("where_filter"): + query_obj = query_obj.with_where(kwargs.get("where_filter")) + query_obj = query_obj.with_additional(["vector"]) + properties = ["text"] + result = query_obj.with_bm25(query=query, properties=properties).with_limit(kwargs.get("top_k", 4)).do() + if "errors" in result: + raise ValueError(f"Error during query: {result['errors']}") + docs = [] + for res in result["data"]["Get"][collection_name]: + text = res.pop(Field.TEXT_KEY.value) + additional = res.pop("_additional") + docs.append(Document(page_content=text, vector=additional["vector"], metadata=res)) + return docs + + def _default_schema(self, index_name: str) -> dict: + return { + "class": index_name, + "properties": [ + { + "name": "text", + "dataType": ["text"], + } + ], + } + + def _json_serializable(self, value: Any) -> Any: + if isinstance(value, datetime.datetime): + return value.isoformat() + return value + + +class WeaviateVectorFactory(AbstractVectorFactory): + def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> WeaviateVector: + if dataset.index_struct_dict: + class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] + collection_name = class_prefix + else: + dataset_id = dataset.id + collection_name = Dataset.gen_collection_name_by_id(dataset_id) + dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.WEAVIATE, collection_name)) + + return WeaviateVector( + collection_name=collection_name, + config=WeaviateConfig( + endpoint=dify_config.WEAVIATE_ENDPOINT or "", + api_key=dify_config.WEAVIATE_API_KEY, + batch_size=dify_config.WEAVIATE_BATCH_SIZE, + ), + attributes=attributes, + ) diff --git a/api/core/rag/docstore/__init__.py b/api/core/rag/docstore/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/docstore/dataset_docstore.py b/api/core/rag/docstore/dataset_docstore.py new file mode 100644 index 0000000000000000000000000000000000000000..8b95d81cc1124b14536aa2e42d7681f0bc601f5c --- /dev/null +++ b/api/core/rag/docstore/dataset_docstore.py @@ -0,0 +1,237 @@ +from collections.abc import Sequence +from typing import Any, Optional + +from sqlalchemy import func + +from core.model_manager import ModelManager +from core.model_runtime.entities.model_entities import ModelType +from core.rag.models.document import Document +from extensions.ext_database import db +from models.dataset import ChildChunk, Dataset, DocumentSegment + + +class DatasetDocumentStore: + def __init__( + self, + dataset: Dataset, + user_id: str, + document_id: Optional[str] = None, + ): + self._dataset = dataset + self._user_id = user_id + self._document_id = document_id + + @classmethod + def from_dict(cls, config_dict: dict[str, Any]) -> "DatasetDocumentStore": + return cls(**config_dict) + + def to_dict(self) -> dict[str, Any]: + """Serialize to dict.""" + return { + "dataset_id": self._dataset.id, + } + + @property + def dateset_id(self) -> Any: + return self._dataset.id + + @property + def user_id(self) -> Any: + return self._user_id + + @property + def docs(self) -> dict[str, Document]: + document_segments = ( + db.session.query(DocumentSegment).filter(DocumentSegment.dataset_id == self._dataset.id).all() + ) + + output = {} + for document_segment in document_segments: + doc_id = document_segment.index_node_id + output[doc_id] = Document( + page_content=document_segment.content, + metadata={ + "doc_id": document_segment.index_node_id, + "doc_hash": document_segment.index_node_hash, + "document_id": document_segment.document_id, + "dataset_id": document_segment.dataset_id, + }, + ) + + return output + + def add_documents(self, docs: Sequence[Document], allow_update: bool = True, save_child: bool = False) -> None: + max_position = ( + db.session.query(func.max(DocumentSegment.position)) + .filter(DocumentSegment.document_id == self._document_id) + .scalar() + ) + + if max_position is None: + max_position = 0 + embedding_model = None + if self._dataset.indexing_technique == "high_quality": + model_manager = ModelManager() + embedding_model = model_manager.get_model_instance( + tenant_id=self._dataset.tenant_id, + provider=self._dataset.embedding_model_provider, + model_type=ModelType.TEXT_EMBEDDING, + model=self._dataset.embedding_model, + ) + + for doc in docs: + if not isinstance(doc, Document): + raise ValueError("doc must be a Document") + + if doc.metadata is None: + raise ValueError("doc.metadata must be a dict") + + segment_document = self.get_document_segment(doc_id=doc.metadata["doc_id"]) + + # NOTE: doc could already exist in the store, but we overwrite it + if not allow_update and segment_document: + raise ValueError( + f"doc_id {doc.metadata['doc_id']} already exists. Set allow_update to True to overwrite." + ) + + # calc embedding use tokens + if embedding_model: + tokens = embedding_model.get_text_embedding_num_tokens(texts=[doc.page_content]) + else: + tokens = 0 + + if not segment_document: + max_position += 1 + + segment_document = DocumentSegment( + tenant_id=self._dataset.tenant_id, + dataset_id=self._dataset.id, + document_id=self._document_id, + index_node_id=doc.metadata["doc_id"], + index_node_hash=doc.metadata["doc_hash"], + position=max_position, + content=doc.page_content, + word_count=len(doc.page_content), + tokens=tokens, + enabled=False, + created_by=self._user_id, + ) + if doc.metadata.get("answer"): + segment_document.answer = doc.metadata.pop("answer", "") + + db.session.add(segment_document) + db.session.flush() + if save_child: + if doc.children: + for postion, child in enumerate(doc.children, start=1): + child_segment = ChildChunk( + tenant_id=self._dataset.tenant_id, + dataset_id=self._dataset.id, + document_id=self._document_id, + segment_id=segment_document.id, + position=postion, + index_node_id=child.metadata.get("doc_id"), + index_node_hash=child.metadata.get("doc_hash"), + content=child.page_content, + word_count=len(child.page_content), + type="automatic", + created_by=self._user_id, + ) + db.session.add(child_segment) + else: + segment_document.content = doc.page_content + if doc.metadata.get("answer"): + segment_document.answer = doc.metadata.pop("answer", "") + segment_document.index_node_hash = doc.metadata.get("doc_hash") + segment_document.word_count = len(doc.page_content) + segment_document.tokens = tokens + if save_child and doc.children: + # delete the existing child chunks + db.session.query(ChildChunk).filter( + ChildChunk.tenant_id == self._dataset.tenant_id, + ChildChunk.dataset_id == self._dataset.id, + ChildChunk.document_id == self._document_id, + ChildChunk.segment_id == segment_document.id, + ).delete() + # add new child chunks + for position, child in enumerate(doc.children, start=1): + child_segment = ChildChunk( + tenant_id=self._dataset.tenant_id, + dataset_id=self._dataset.id, + document_id=self._document_id, + segment_id=segment_document.id, + position=position, + index_node_id=child.metadata.get("doc_id"), + index_node_hash=child.metadata.get("doc_hash"), + content=child.page_content, + word_count=len(child.page_content), + type="automatic", + created_by=self._user_id, + ) + db.session.add(child_segment) + + db.session.commit() + + def document_exists(self, doc_id: str) -> bool: + """Check if document exists.""" + result = self.get_document_segment(doc_id) + return result is not None + + def get_document(self, doc_id: str, raise_error: bool = True) -> Optional[Document]: + document_segment = self.get_document_segment(doc_id) + + if document_segment is None: + if raise_error: + raise ValueError(f"doc_id {doc_id} not found.") + else: + return None + + return Document( + page_content=document_segment.content, + metadata={ + "doc_id": document_segment.index_node_id, + "doc_hash": document_segment.index_node_hash, + "document_id": document_segment.document_id, + "dataset_id": document_segment.dataset_id, + }, + ) + + def delete_document(self, doc_id: str, raise_error: bool = True) -> None: + document_segment = self.get_document_segment(doc_id) + + if document_segment is None: + if raise_error: + raise ValueError(f"doc_id {doc_id} not found.") + else: + return None + + db.session.delete(document_segment) + db.session.commit() + + def set_document_hash(self, doc_id: str, doc_hash: str) -> None: + """Set the hash for a given doc_id.""" + document_segment = self.get_document_segment(doc_id) + + if document_segment is None: + return None + + document_segment.index_node_hash = doc_hash + db.session.commit() + + def get_document_hash(self, doc_id: str) -> Optional[str]: + """Get the stored hash for a document, if it exists.""" + document_segment = self.get_document_segment(doc_id) + + if document_segment is None: + return None + data: Optional[str] = document_segment.index_node_hash + return data + + def get_document_segment(self, doc_id: str) -> Optional[DocumentSegment]: + document_segment = ( + db.session.query(DocumentSegment) + .filter(DocumentSegment.dataset_id == self._dataset.id, DocumentSegment.index_node_id == doc_id) + .first() + ) + + return document_segment diff --git a/api/core/rag/embedding/__init__.py b/api/core/rag/embedding/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/embedding/cached_embedding.py b/api/core/rag/embedding/cached_embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..a2c8737da791984afac6ae9ad262bf892ef4a098 --- /dev/null +++ b/api/core/rag/embedding/cached_embedding.py @@ -0,0 +1,142 @@ +import base64 +import logging +from typing import Any, Optional, cast + +import numpy as np +from sqlalchemy.exc import IntegrityError + +from configs import dify_config +from core.entities.embedding_type import EmbeddingInputType +from core.model_manager import ModelInstance +from core.model_runtime.entities.model_entities import ModelPropertyKey +from core.model_runtime.model_providers.__base.text_embedding_model import TextEmbeddingModel +from core.rag.embedding.embedding_base import Embeddings +from extensions.ext_database import db +from extensions.ext_redis import redis_client +from libs import helper +from models.dataset import Embedding + +logger = logging.getLogger(__name__) + + +class CacheEmbedding(Embeddings): + def __init__(self, model_instance: ModelInstance, user: Optional[str] = None) -> None: + self._model_instance = model_instance + self._user = user + + def embed_documents(self, texts: list[str]) -> list[list[float]]: + """Embed search docs in batches of 10.""" + # use doc embedding cache or store if not exists + text_embeddings: list[Any] = [None for _ in range(len(texts))] + embedding_queue_indices = [] + for i, text in enumerate(texts): + hash = helper.generate_text_hash(text) + embedding = ( + db.session.query(Embedding) + .filter_by( + model_name=self._model_instance.model, hash=hash, provider_name=self._model_instance.provider + ) + .first() + ) + if embedding: + text_embeddings[i] = embedding.get_embedding() + else: + embedding_queue_indices.append(i) + if embedding_queue_indices: + embedding_queue_texts = [texts[i] for i in embedding_queue_indices] + embedding_queue_embeddings = [] + try: + model_type_instance = cast(TextEmbeddingModel, self._model_instance.model_type_instance) + model_schema = model_type_instance.get_model_schema( + self._model_instance.model, self._model_instance.credentials + ) + max_chunks = ( + model_schema.model_properties[ModelPropertyKey.MAX_CHUNKS] + if model_schema and ModelPropertyKey.MAX_CHUNKS in model_schema.model_properties + else 1 + ) + for i in range(0, len(embedding_queue_texts), max_chunks): + batch_texts = embedding_queue_texts[i : i + max_chunks] + + embedding_result = self._model_instance.invoke_text_embedding( + texts=batch_texts, user=self._user, input_type=EmbeddingInputType.DOCUMENT + ) + + for vector in embedding_result.embeddings: + try: + # FIXME: type ignore for numpy here + normalized_embedding = (vector / np.linalg.norm(vector)).tolist() # type: ignore + # stackoverflow best way: https://stackoverflow.com/questions/20319813/how-to-check-list-containing-nan + if np.isnan(normalized_embedding).any(): + # for issue #11827 float values are not json compliant + logger.warning(f"Normalized embedding is nan: {normalized_embedding}") + continue + embedding_queue_embeddings.append(normalized_embedding) + except IntegrityError: + db.session.rollback() + except Exception as e: + logging.exception("Failed transform embedding") + cache_embeddings = [] + try: + for i, n_embedding in zip(embedding_queue_indices, embedding_queue_embeddings): + text_embeddings[i] = n_embedding + hash = helper.generate_text_hash(texts[i]) + if hash not in cache_embeddings: + embedding_cache = Embedding( + model_name=self._model_instance.model, + hash=hash, + provider_name=self._model_instance.provider, + ) + embedding_cache.set_embedding(n_embedding) + db.session.add(embedding_cache) + cache_embeddings.append(hash) + db.session.commit() + except IntegrityError: + db.session.rollback() + except Exception as ex: + db.session.rollback() + logger.exception("Failed to embed documents: %s") + raise ex + + return text_embeddings + + def embed_query(self, text: str) -> list[float]: + """Embed query text.""" + # use doc embedding cache or store if not exists + hash = helper.generate_text_hash(text) + embedding_cache_key = f"{self._model_instance.provider}_{self._model_instance.model}_{hash}" + embedding = redis_client.get(embedding_cache_key) + if embedding: + redis_client.expire(embedding_cache_key, 600) + decoded_embedding = np.frombuffer(base64.b64decode(embedding), dtype="float") + return [float(x) for x in decoded_embedding] + try: + embedding_result = self._model_instance.invoke_text_embedding( + texts=[text], user=self._user, input_type=EmbeddingInputType.QUERY + ) + + embedding_results = embedding_result.embeddings[0] + # FIXME: type ignore for numpy here + embedding_results = (embedding_results / np.linalg.norm(embedding_results)).tolist() # type: ignore + if np.isnan(embedding_results).any(): + raise ValueError("Normalized embedding is nan please try again") + except Exception as ex: + if dify_config.DEBUG: + logging.exception(f"Failed to embed query text '{text[:10]}...({len(text)} chars)'") + raise ex + + try: + # encode embedding to base64 + embedding_vector = np.array(embedding_results) + vector_bytes = embedding_vector.tobytes() + # Transform to Base64 + encoded_vector = base64.b64encode(vector_bytes) + # Transform to string + encoded_str = encoded_vector.decode("utf-8") + redis_client.setex(embedding_cache_key, 600, encoded_str) + except Exception as ex: + if dify_config.DEBUG: + logging.exception(f"Failed to add embedding to redis for the text '{text[:10]}...({len(text)} chars)'") + raise ex + + return embedding_results diff --git a/api/core/rag/embedding/embedding_base.py b/api/core/rag/embedding/embedding_base.py new file mode 100644 index 0000000000000000000000000000000000000000..9f232ab91089fec63d916ff6bac51002d92cc73d --- /dev/null +++ b/api/core/rag/embedding/embedding_base.py @@ -0,0 +1,23 @@ +from abc import ABC, abstractmethod + + +class Embeddings(ABC): + """Interface for embedding models.""" + + @abstractmethod + def embed_documents(self, texts: list[str]) -> list[list[float]]: + """Embed search docs.""" + raise NotImplementedError + + @abstractmethod + def embed_query(self, text: str) -> list[float]: + """Embed query text.""" + raise NotImplementedError + + async def aembed_documents(self, texts: list[str]) -> list[list[float]]: + """Asynchronous Embed search docs.""" + raise NotImplementedError + + async def aembed_query(self, text: str) -> list[float]: + """Asynchronous Embed query text.""" + raise NotImplementedError diff --git a/api/core/rag/embedding/retrieval.py b/api/core/rag/embedding/retrieval.py new file mode 100644 index 0000000000000000000000000000000000000000..800422d888e4fad24f541b4084f76d1a6396fe35 --- /dev/null +++ b/api/core/rag/embedding/retrieval.py @@ -0,0 +1,23 @@ +from typing import Optional + +from pydantic import BaseModel + +from models.dataset import DocumentSegment + + +class RetrievalChildChunk(BaseModel): + """Retrieval segments.""" + + id: str + content: str + score: float + position: int + + +class RetrievalSegments(BaseModel): + """Retrieval segments.""" + + model_config = {"arbitrary_types_allowed": True} + segment: DocumentSegment + child_chunks: Optional[list[RetrievalChildChunk]] = None + score: Optional[float] = None diff --git a/api/core/rag/entities/context_entities.py b/api/core/rag/entities/context_entities.py new file mode 100644 index 0000000000000000000000000000000000000000..cd18ad081ff4fd772203a7c8421c53ee7faab69a --- /dev/null +++ b/api/core/rag/entities/context_entities.py @@ -0,0 +1,12 @@ +from typing import Optional + +from pydantic import BaseModel + + +class DocumentContext(BaseModel): + """ + Model class for document context. + """ + + content: str + score: Optional[float] = None diff --git a/api/core/rag/extractor/blob/blob.py b/api/core/rag/extractor/blob/blob.py new file mode 100644 index 0000000000000000000000000000000000000000..e46ab8b7fd0ac2a1a0d879e916fdd59d85ea229d --- /dev/null +++ b/api/core/rag/extractor/blob/blob.py @@ -0,0 +1,163 @@ +"""Schema for Blobs and Blob Loaders. + +The goal is to facilitate decoupling of content loading from content parsing code. + +In addition, content loading code should provide a lazy loading interface by default. +""" + +from __future__ import annotations + +import contextlib +import mimetypes +from abc import ABC, abstractmethod +from collections.abc import Generator, Iterable, Mapping +from io import BufferedReader, BytesIO +from pathlib import Path, PurePath +from typing import Any, Optional, Union + +from pydantic import BaseModel, ConfigDict, model_validator + +PathLike = Union[str, PurePath] + + +class Blob(BaseModel): + """A blob is used to represent raw data by either reference or value. + + Provides an interface to materialize the blob in different representations, and + help to decouple the development of data loaders from the downstream parsing of + the raw data. + + Inspired by: https://developer.mozilla.org/en-US/docs/Web/API/Blob + """ + + data: Union[bytes, str, None] = None # Raw data + mimetype: Optional[str] = None # Not to be confused with a file extension + encoding: str = "utf-8" # Use utf-8 as default encoding, if decoding to string + # Location where the original content was found + # Represent location on the local file system + # Useful for situations where downstream code assumes it must work with file paths + # rather than in-memory content. + path: Optional[PathLike] = None + model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True) + + @property + def source(self) -> Optional[str]: + """The source location of the blob as string if known otherwise none.""" + return str(self.path) if self.path else None + + @model_validator(mode="before") + @classmethod + def check_blob_is_valid(cls, values: Mapping[str, Any]) -> Mapping[str, Any]: + """Verify that either data or path is provided.""" + if "data" not in values and "path" not in values: + raise ValueError("Either data or path must be provided") + return values + + def as_string(self) -> str: + """Read data as a string.""" + if self.data is None and self.path: + return Path(str(self.path)).read_text(encoding=self.encoding) + elif isinstance(self.data, bytes): + return self.data.decode(self.encoding) + elif isinstance(self.data, str): + return self.data + else: + raise ValueError(f"Unable to get string for blob {self}") + + def as_bytes(self) -> bytes: + """Read data as bytes.""" + if isinstance(self.data, bytes): + return self.data + elif isinstance(self.data, str): + return self.data.encode(self.encoding) + elif self.data is None and self.path: + return Path(str(self.path)).read_bytes() + else: + raise ValueError(f"Unable to get bytes for blob {self}") + + @contextlib.contextmanager + def as_bytes_io(self) -> Generator[Union[BytesIO, BufferedReader], None, None]: + """Read data as a byte stream.""" + if isinstance(self.data, bytes): + yield BytesIO(self.data) + elif self.data is None and self.path: + with open(str(self.path), "rb") as f: + yield f + else: + raise NotImplementedError(f"Unable to convert blob {self}") + + @classmethod + def from_path( + cls, + path: PathLike, + *, + encoding: str = "utf-8", + mime_type: Optional[str] = None, + guess_type: bool = True, + ) -> Blob: + """Load the blob from a path like object. + + Args: + path: path like object to file to be read + encoding: Encoding to use if decoding the bytes into a string + mime_type: if provided, will be set as the mime-type of the data + guess_type: If True, the mimetype will be guessed from the file extension, + if a mime-type was not provided + + Returns: + Blob instance + """ + if mime_type is None and guess_type: + _mimetype = mimetypes.guess_type(path)[0] if guess_type else None + else: + _mimetype = mime_type + # We do not load the data immediately, instead we treat the blob as a + # reference to the underlying data. + return cls(data=None, mimetype=_mimetype, encoding=encoding, path=path) + + @classmethod + def from_data( + cls, + data: Union[str, bytes], + *, + encoding: str = "utf-8", + mime_type: Optional[str] = None, + path: Optional[str] = None, + ) -> Blob: + """Initialize the blob from in-memory data. + + Args: + data: the in-memory data associated with the blob + encoding: Encoding to use if decoding the bytes into a string + mime_type: if provided, will be set as the mime-type of the data + path: if provided, will be set as the source from which the data came + + Returns: + Blob instance + """ + return cls(data=data, mimetype=mime_type, encoding=encoding, path=path) + + def __repr__(self) -> str: + """Define the blob representation.""" + str_repr = f"Blob {id(self)}" + if self.source: + str_repr += f" {self.source}" + return str_repr + + +class BlobLoader(ABC): + """Abstract interface for blob loaders implementation. + + Implementer should be able to load raw content from a datasource system according + to some criteria and return the raw content lazily as a stream of blobs. + """ + + @abstractmethod + def yield_blobs( + self, + ) -> Iterable[Blob]: + """A lazy loader for raw data represented by Blob object. + + Returns: + A generator over blobs + """ diff --git a/api/core/rag/extractor/csv_extractor.py b/api/core/rag/extractor/csv_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..5b674039024189c644ff2c896c18cbb079c6828d --- /dev/null +++ b/api/core/rag/extractor/csv_extractor.py @@ -0,0 +1,78 @@ +"""Abstract interface for document loader implementations.""" + +import csv +from typing import Optional + +import pandas as pd + +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.extractor.helpers import detect_file_encodings +from core.rag.models.document import Document + + +class CSVExtractor(BaseExtractor): + """Load CSV files. + + + Args: + file_path: Path to the file to load. + """ + + def __init__( + self, + file_path: str, + encoding: Optional[str] = None, + autodetect_encoding: bool = False, + source_column: Optional[str] = None, + csv_args: Optional[dict] = None, + ): + """Initialize with file path.""" + self._file_path = file_path + self._encoding = encoding + self._autodetect_encoding = autodetect_encoding + self.source_column = source_column + self.csv_args = csv_args or {} + + def extract(self) -> list[Document]: + """Load data into document objects.""" + docs = [] + try: + with open(self._file_path, newline="", encoding=self._encoding) as csvfile: + docs = self._read_from_file(csvfile) + except UnicodeDecodeError as e: + if self._autodetect_encoding: + detected_encodings = detect_file_encodings(self._file_path) + for encoding in detected_encodings: + try: + with open(self._file_path, newline="", encoding=encoding.encoding) as csvfile: + docs = self._read_from_file(csvfile) + break + except UnicodeDecodeError: + continue + else: + raise RuntimeError(f"Error loading {self._file_path}") from e + + return docs + + def _read_from_file(self, csvfile) -> list[Document]: + docs = [] + try: + # load csv file into pandas dataframe + df = pd.read_csv(csvfile, on_bad_lines="skip", **self.csv_args) + + # check source column exists + if self.source_column and self.source_column not in df.columns: + raise ValueError(f"Source column '{self.source_column}' not found in CSV file.") + + # create document objects + + for i, row in df.iterrows(): + content = ";".join(f"{col.strip()}: {str(row[col]).strip()}" for col in df.columns) + source = row[self.source_column] if self.source_column else "" + metadata = {"source": source, "row": i} + doc = Document(page_content=content, metadata=metadata) + docs.append(doc) + except csv.Error as e: + raise e + + return docs diff --git a/api/core/rag/extractor/entity/datasource_type.py b/api/core/rag/extractor/entity/datasource_type.py new file mode 100644 index 0000000000000000000000000000000000000000..19ad300d110fe63ea3891db5d62b5aae7b457ee9 --- /dev/null +++ b/api/core/rag/extractor/entity/datasource_type.py @@ -0,0 +1,7 @@ +from enum import Enum + + +class DatasourceType(Enum): + FILE = "upload_file" + NOTION = "notion_import" + WEBSITE = "website_crawl" diff --git a/api/core/rag/extractor/entity/extract_setting.py b/api/core/rag/extractor/entity/extract_setting.py new file mode 100644 index 0000000000000000000000000000000000000000..7c00c668dd49a3309184c21000f3ab9c40aa4cba --- /dev/null +++ b/api/core/rag/extractor/entity/extract_setting.py @@ -0,0 +1,57 @@ +from typing import Optional + +from pydantic import BaseModel, ConfigDict + +from models.dataset import Document +from models.model import UploadFile + + +class NotionInfo(BaseModel): + """ + Notion import info. + """ + + notion_workspace_id: str + notion_obj_id: str + notion_page_type: str + document: Optional[Document] = None + tenant_id: str + model_config = ConfigDict(arbitrary_types_allowed=True) + + def __init__(self, **data) -> None: + super().__init__(**data) + + +class WebsiteInfo(BaseModel): + """ + website import info. + """ + + provider: str + job_id: str + url: str + mode: str + tenant_id: str + only_main_content: bool = False + + class Config: + arbitrary_types_allowed = True + + def __init__(self, **data) -> None: + super().__init__(**data) + + +class ExtractSetting(BaseModel): + """ + Model class for provider response. + """ + + datasource_type: str + upload_file: Optional[UploadFile] = None + notion_info: Optional[NotionInfo] = None + website_info: Optional[WebsiteInfo] = None + document_model: Optional[str] = None + model_config = ConfigDict(arbitrary_types_allowed=True) + + def __init__(self, **data) -> None: + super().__init__(**data) diff --git a/api/core/rag/extractor/excel_extractor.py b/api/core/rag/extractor/excel_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..a3b35458df9ab0d020e62fb13adb1b00b7c9afb4 --- /dev/null +++ b/api/core/rag/extractor/excel_extractor.py @@ -0,0 +1,78 @@ +"""Abstract interface for document loader implementations.""" + +import os +from typing import Optional, cast + +import pandas as pd +from openpyxl import load_workbook # type: ignore + +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.models.document import Document + + +class ExcelExtractor(BaseExtractor): + """Load Excel files. + + + Args: + file_path: Path to the file to load. + """ + + def __init__(self, file_path: str, encoding: Optional[str] = None, autodetect_encoding: bool = False): + """Initialize with file path.""" + self._file_path = file_path + self._encoding = encoding + self._autodetect_encoding = autodetect_encoding + + def extract(self) -> list[Document]: + """Load from Excel file in xls or xlsx format using Pandas and openpyxl.""" + documents = [] + file_extension = os.path.splitext(self._file_path)[-1].lower() + + if file_extension == ".xlsx": + wb = load_workbook(self._file_path, data_only=True) + for sheet_name in wb.sheetnames: + sheet = wb[sheet_name] + data = sheet.values + try: + cols = next(data) + except StopIteration: + continue + df = pd.DataFrame(data, columns=cols) + + df.dropna(how="all", inplace=True) + + for index, row in df.iterrows(): + page_content = [] + for col_index, (k, v) in enumerate(row.items()): + if pd.notna(v): + cell = sheet.cell( + row=cast(int, index) + 2, column=col_index + 1 + ) # +2 to account for header and 1-based index + if cell.hyperlink: + value = f"[{v}]({cell.hyperlink.target})" + page_content.append(f'"{k}":"{value}"') + else: + page_content.append(f'"{k}":"{v}"') + documents.append( + Document(page_content=";".join(page_content), metadata={"source": self._file_path}) + ) + + elif file_extension == ".xls": + excel_file = pd.ExcelFile(self._file_path, engine="xlrd") + for excel_sheet_name in excel_file.sheet_names: + df = excel_file.parse(sheet_name=excel_sheet_name) + df.dropna(how="all", inplace=True) + + for _, row in df.iterrows(): + page_content = [] + for k, v in row.items(): + if pd.notna(v): + page_content.append(f'"{k}":"{v}"') + documents.append( + Document(page_content=";".join(page_content), metadata={"source": self._file_path}) + ) + else: + raise ValueError(f"Unsupported file extension: {file_extension}") + + return documents diff --git a/api/core/rag/extractor/extract_processor.py b/api/core/rag/extractor/extract_processor.py new file mode 100644 index 0000000000000000000000000000000000000000..f9fd7f92a12c81015d0542204d2095a80371dda6 --- /dev/null +++ b/api/core/rag/extractor/extract_processor.py @@ -0,0 +1,195 @@ +import re +import tempfile +from pathlib import Path +from typing import Optional, Union +from urllib.parse import unquote + +from configs import dify_config +from core.helper import ssrf_proxy +from core.rag.extractor.csv_extractor import CSVExtractor +from core.rag.extractor.entity.datasource_type import DatasourceType +from core.rag.extractor.entity.extract_setting import ExtractSetting +from core.rag.extractor.excel_extractor import ExcelExtractor +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.extractor.firecrawl.firecrawl_web_extractor import FirecrawlWebExtractor +from core.rag.extractor.html_extractor import HtmlExtractor +from core.rag.extractor.jina_reader_extractor import JinaReaderWebExtractor +from core.rag.extractor.markdown_extractor import MarkdownExtractor +from core.rag.extractor.notion_extractor import NotionExtractor +from core.rag.extractor.pdf_extractor import PdfExtractor +from core.rag.extractor.text_extractor import TextExtractor +from core.rag.extractor.unstructured.unstructured_eml_extractor import UnstructuredEmailExtractor +from core.rag.extractor.unstructured.unstructured_epub_extractor import UnstructuredEpubExtractor +from core.rag.extractor.unstructured.unstructured_markdown_extractor import UnstructuredMarkdownExtractor +from core.rag.extractor.unstructured.unstructured_msg_extractor import UnstructuredMsgExtractor +from core.rag.extractor.unstructured.unstructured_ppt_extractor import UnstructuredPPTExtractor +from core.rag.extractor.unstructured.unstructured_pptx_extractor import UnstructuredPPTXExtractor +from core.rag.extractor.unstructured.unstructured_xml_extractor import UnstructuredXmlExtractor +from core.rag.extractor.word_extractor import WordExtractor +from core.rag.models.document import Document +from extensions.ext_storage import storage +from models.model import UploadFile + +SUPPORT_URL_CONTENT_TYPES = ["application/pdf", "text/plain", "application/json"] +USER_AGENT = ( + "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124" + " Safari/537.36" +) + + +class ExtractProcessor: + @classmethod + def load_from_upload_file( + cls, upload_file: UploadFile, return_text: bool = False, is_automatic: bool = False + ) -> Union[list[Document], str]: + extract_setting = ExtractSetting( + datasource_type="upload_file", upload_file=upload_file, document_model="text_model" + ) + if return_text: + delimiter = "\n" + return delimiter.join([document.page_content for document in cls.extract(extract_setting, is_automatic)]) + else: + return cls.extract(extract_setting, is_automatic) + + @classmethod + def load_from_url(cls, url: str, return_text: bool = False) -> Union[list[Document], str]: + response = ssrf_proxy.get(url, headers={"User-Agent": USER_AGENT}) + + with tempfile.TemporaryDirectory() as temp_dir: + suffix = Path(url).suffix + if not suffix and suffix != ".": + # get content-type + if response.headers.get("Content-Type"): + suffix = "." + response.headers.get("Content-Type").split("/")[-1] + else: + content_disposition = response.headers.get("Content-Disposition") + filename_match = re.search(r'filename="([^"]+)"', content_disposition) + if filename_match: + filename = unquote(filename_match.group(1)) + match = re.search(r"\.(\w+)$", filename) + if match: + suffix = "." + match.group(1) + else: + suffix = "" + # FIXME mypy: Cannot determine type of 'tempfile._get_candidate_names' better not use it here + file_path = f"{temp_dir}/{next(tempfile._get_candidate_names())}{suffix}" # type: ignore + Path(file_path).write_bytes(response.content) + extract_setting = ExtractSetting(datasource_type="upload_file", document_model="text_model") + if return_text: + delimiter = "\n" + return delimiter.join( + [ + document.page_content + for document in cls.extract(extract_setting=extract_setting, file_path=file_path) + ] + ) + else: + return cls.extract(extract_setting=extract_setting, file_path=file_path) + + @classmethod + def extract( + cls, extract_setting: ExtractSetting, is_automatic: bool = False, file_path: Optional[str] = None + ) -> list[Document]: + if extract_setting.datasource_type == DatasourceType.FILE.value: + with tempfile.TemporaryDirectory() as temp_dir: + if not file_path: + assert extract_setting.upload_file is not None, "upload_file is required" + upload_file: UploadFile = extract_setting.upload_file + suffix = Path(upload_file.key).suffix + # FIXME mypy: Cannot determine type of 'tempfile._get_candidate_names' better not use it here + file_path = f"{temp_dir}/{next(tempfile._get_candidate_names())}{suffix}" # type: ignore + storage.download(upload_file.key, file_path) + input_file = Path(file_path) + file_extension = input_file.suffix.lower() + etl_type = dify_config.ETL_TYPE + extractor: Optional[BaseExtractor] = None + if etl_type == "Unstructured": + unstructured_api_url = dify_config.UNSTRUCTURED_API_URL + unstructured_api_key = dify_config.UNSTRUCTURED_API_KEY or "" + + if file_extension in {".xlsx", ".xls"}: + extractor = ExcelExtractor(file_path) + elif file_extension == ".pdf": + extractor = PdfExtractor(file_path) + elif file_extension in {".md", ".markdown", ".mdx"}: + extractor = ( + UnstructuredMarkdownExtractor(file_path, unstructured_api_url, unstructured_api_key) + if is_automatic + else MarkdownExtractor(file_path, autodetect_encoding=True) + ) + elif file_extension in {".htm", ".html"}: + extractor = HtmlExtractor(file_path) + elif file_extension == ".docx": + extractor = WordExtractor(file_path, upload_file.tenant_id, upload_file.created_by) + elif file_extension == ".csv": + extractor = CSVExtractor(file_path, autodetect_encoding=True) + elif file_extension == ".msg": + extractor = UnstructuredMsgExtractor(file_path, unstructured_api_url, unstructured_api_key) + elif file_extension == ".eml": + extractor = UnstructuredEmailExtractor(file_path, unstructured_api_url, unstructured_api_key) + elif file_extension == ".ppt": + extractor = UnstructuredPPTExtractor(file_path, unstructured_api_url, unstructured_api_key) + # You must first specify the API key + # because unstructured_api_key is necessary to parse .ppt documents + elif file_extension == ".pptx": + extractor = UnstructuredPPTXExtractor(file_path, unstructured_api_url, unstructured_api_key) + elif file_extension == ".xml": + extractor = UnstructuredXmlExtractor(file_path, unstructured_api_url, unstructured_api_key) + elif file_extension == ".epub": + extractor = UnstructuredEpubExtractor(file_path, unstructured_api_url, unstructured_api_key) + else: + # txt + extractor = TextExtractor(file_path, autodetect_encoding=True) + else: + if file_extension in {".xlsx", ".xls"}: + extractor = ExcelExtractor(file_path) + elif file_extension == ".pdf": + extractor = PdfExtractor(file_path) + elif file_extension in {".md", ".markdown", ".mdx"}: + extractor = MarkdownExtractor(file_path, autodetect_encoding=True) + elif file_extension in {".htm", ".html"}: + extractor = HtmlExtractor(file_path) + elif file_extension == ".docx": + extractor = WordExtractor(file_path, upload_file.tenant_id, upload_file.created_by) + elif file_extension == ".csv": + extractor = CSVExtractor(file_path, autodetect_encoding=True) + elif file_extension == ".epub": + extractor = UnstructuredEpubExtractor(file_path) + else: + # txt + extractor = TextExtractor(file_path, autodetect_encoding=True) + return extractor.extract() + elif extract_setting.datasource_type == DatasourceType.NOTION.value: + assert extract_setting.notion_info is not None, "notion_info is required" + extractor = NotionExtractor( + notion_workspace_id=extract_setting.notion_info.notion_workspace_id, + notion_obj_id=extract_setting.notion_info.notion_obj_id, + notion_page_type=extract_setting.notion_info.notion_page_type, + document_model=extract_setting.notion_info.document, + tenant_id=extract_setting.notion_info.tenant_id, + ) + return extractor.extract() + elif extract_setting.datasource_type == DatasourceType.WEBSITE.value: + assert extract_setting.website_info is not None, "website_info is required" + if extract_setting.website_info.provider == "firecrawl": + extractor = FirecrawlWebExtractor( + url=extract_setting.website_info.url, + job_id=extract_setting.website_info.job_id, + tenant_id=extract_setting.website_info.tenant_id, + mode=extract_setting.website_info.mode, + only_main_content=extract_setting.website_info.only_main_content, + ) + return extractor.extract() + elif extract_setting.website_info.provider == "jinareader": + extractor = JinaReaderWebExtractor( + url=extract_setting.website_info.url, + job_id=extract_setting.website_info.job_id, + tenant_id=extract_setting.website_info.tenant_id, + mode=extract_setting.website_info.mode, + only_main_content=extract_setting.website_info.only_main_content, + ) + return extractor.extract() + else: + raise ValueError(f"Unsupported website provider: {extract_setting.website_info.provider}") + else: + raise ValueError(f"Unsupported datasource type: {extract_setting.datasource_type}") diff --git a/api/core/rag/extractor/extractor_base.py b/api/core/rag/extractor/extractor_base.py new file mode 100644 index 0000000000000000000000000000000000000000..582eca94df71e177af2ce1870da0a1bcc22af6b8 --- /dev/null +++ b/api/core/rag/extractor/extractor_base.py @@ -0,0 +1,11 @@ +"""Abstract interface for document loader implementations.""" + +from abc import ABC, abstractmethod + + +class BaseExtractor(ABC): + """Interface for extract files.""" + + @abstractmethod + def extract(self): + raise NotImplementedError diff --git a/api/core/rag/extractor/firecrawl/firecrawl_app.py b/api/core/rag/extractor/firecrawl/firecrawl_app.py new file mode 100644 index 0000000000000000000000000000000000000000..836a1398bfdad9a06609afba8345a2fc7a300b3e --- /dev/null +++ b/api/core/rag/extractor/firecrawl/firecrawl_app.py @@ -0,0 +1,129 @@ +import json +import time +from typing import Any, cast + +import requests + +from extensions.ext_storage import storage + + +class FirecrawlApp: + def __init__(self, api_key=None, base_url=None): + self.api_key = api_key + self.base_url = base_url or "https://api.firecrawl.dev" + if self.api_key is None and self.base_url == "https://api.firecrawl.dev": + raise ValueError("No API key provided") + + def scrape_url(self, url, params=None) -> dict[str, Any]: + # Documentation: https://docs.firecrawl.dev/api-reference/endpoint/scrape + headers = self._prepare_headers() + json_data = { + "url": url, + "formats": ["markdown"], + "onlyMainContent": True, + "timeout": 30000, + } + if params: + json_data.update(params) + response = self._post_request(f"{self.base_url}/v1/scrape", json_data, headers) + if response.status_code == 200: + response_data = response.json() + data = response_data["data"] + return self._extract_common_fields(data) + elif response.status_code in {402, 409, 500, 429, 408}: + self._handle_error(response, "scrape URL") + return {} # Avoid additional exception after handling error + else: + raise Exception(f"Failed to scrape URL. Status code: {response.status_code}") + + def crawl_url(self, url, params=None) -> str: + # Documentation: https://docs.firecrawl.dev/api-reference/endpoint/crawl-post + headers = self._prepare_headers() + json_data = {"url": url} + if params: + json_data.update(params) + response = self._post_request(f"{self.base_url}/v1/crawl", json_data, headers) + if response.status_code == 200: + # There's also another two fields in the response: "success" (bool) and "url" (str) + job_id = response.json().get("id") + return cast(str, job_id) + else: + self._handle_error(response, "start crawl job") + # FIXME: unreachable code for mypy + return "" # unreachable + + def check_crawl_status(self, job_id) -> dict[str, Any]: + headers = self._prepare_headers() + response = self._get_request(f"{self.base_url}/v1/crawl/{job_id}", headers) + if response.status_code == 200: + crawl_status_response = response.json() + if crawl_status_response.get("status") == "completed": + total = crawl_status_response.get("total", 0) + if total == 0: + raise Exception("Failed to check crawl status. Error: No page found") + data = crawl_status_response.get("data", []) + url_data_list = [] + for item in data: + if isinstance(item, dict) and "metadata" in item and "markdown" in item: + url_data = self._extract_common_fields(item) + url_data_list.append(url_data) + if url_data_list: + file_key = "website_files/" + job_id + ".txt" + try: + if storage.exists(file_key): + storage.delete(file_key) + storage.save(file_key, json.dumps(url_data_list).encode("utf-8")) + except Exception as e: + raise Exception(f"Error saving crawl data: {e}") + return self._format_crawl_status_response("completed", crawl_status_response, url_data_list) + else: + return self._format_crawl_status_response( + crawl_status_response.get("status"), crawl_status_response, [] + ) + else: + self._handle_error(response, "check crawl status") + # FIXME: unreachable code for mypy + return {} # unreachable + + def _format_crawl_status_response( + self, status: str, crawl_status_response: dict[str, Any], url_data_list: list[dict[str, Any]] + ) -> dict[str, Any]: + return { + "status": status, + "total": crawl_status_response.get("total"), + "current": crawl_status_response.get("completed"), + "data": url_data_list, + } + + def _extract_common_fields(self, item: dict[str, Any]) -> dict[str, Any]: + return { + "title": item.get("metadata", {}).get("title"), + "description": item.get("metadata", {}).get("description"), + "source_url": item.get("metadata", {}).get("sourceURL"), + "markdown": item.get("markdown"), + } + + def _prepare_headers(self) -> dict[str, Any]: + return {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"} + + def _post_request(self, url, data, headers, retries=3, backoff_factor=0.5) -> requests.Response: + for attempt in range(retries): + response = requests.post(url, headers=headers, json=data) + if response.status_code == 502: + time.sleep(backoff_factor * (2**attempt)) + else: + return response + return response + + def _get_request(self, url, headers, retries=3, backoff_factor=0.5) -> requests.Response: + for attempt in range(retries): + response = requests.get(url, headers=headers) + if response.status_code == 502: + time.sleep(backoff_factor * (2**attempt)) + else: + return response + return response + + def _handle_error(self, response, action) -> None: + error_message = response.json().get("error", "Unknown error occurred") + raise Exception(f"Failed to {action}. Status code: {response.status_code}. Error: {error_message}") diff --git a/api/core/rag/extractor/firecrawl/firecrawl_web_extractor.py b/api/core/rag/extractor/firecrawl/firecrawl_web_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..355a2fb2048983e68e58176675e9406fefdfab8c --- /dev/null +++ b/api/core/rag/extractor/firecrawl/firecrawl_web_extractor.py @@ -0,0 +1,57 @@ +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.models.document import Document +from services.website_service import WebsiteService + + +class FirecrawlWebExtractor(BaseExtractor): + """ + Crawl and scrape websites and return content in clean llm-ready markdown. + + + Args: + url: The URL to scrape. + api_key: The API key for Firecrawl. + base_url: The base URL for the Firecrawl API. Defaults to 'https://api.firecrawl.dev'. + mode: The mode of operation. Defaults to 'scrape'. Options are 'crawl', 'scrape' and 'crawl_return_urls'. + only_main_content: Only return the main content of the page excluding headers, navs, footers, etc. + """ + + def __init__(self, url: str, job_id: str, tenant_id: str, mode: str = "crawl", only_main_content: bool = True): + """Initialize with url, api_key, base_url and mode.""" + self._url = url + self.job_id = job_id + self.tenant_id = tenant_id + self.mode = mode + self.only_main_content = only_main_content + + def extract(self) -> list[Document]: + """Extract content from the URL.""" + documents = [] + if self.mode == "crawl": + crawl_data = WebsiteService.get_crawl_url_data(self.job_id, "firecrawl", self._url, self.tenant_id) + if crawl_data is None: + return [] + document = Document( + page_content=crawl_data.get("markdown", ""), + metadata={ + "source_url": crawl_data.get("source_url"), + "description": crawl_data.get("description"), + "title": crawl_data.get("title"), + }, + ) + documents.append(document) + elif self.mode == "scrape": + scrape_data = WebsiteService.get_scrape_url_data( + "firecrawl", self._url, self.tenant_id, self.only_main_content + ) + + document = Document( + page_content=scrape_data.get("markdown", ""), + metadata={ + "source_url": scrape_data.get("source_url"), + "description": scrape_data.get("description"), + "title": scrape_data.get("title"), + }, + ) + documents.append(document) + return documents diff --git a/api/core/rag/extractor/helpers.py b/api/core/rag/extractor/helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..69ca9d5d63688834a3322e65d9fdacb495422aab --- /dev/null +++ b/api/core/rag/extractor/helpers.py @@ -0,0 +1,44 @@ +"""Document loader helpers.""" + +import concurrent.futures +from pathlib import Path +from typing import NamedTuple, Optional, cast + + +class FileEncoding(NamedTuple): + """A file encoding as the NamedTuple.""" + + encoding: Optional[str] + """The encoding of the file.""" + confidence: float + """The confidence of the encoding.""" + language: Optional[str] + """The language of the file.""" + + +def detect_file_encodings(file_path: str, timeout: int = 5) -> list[FileEncoding]: + """Try to detect the file encoding. + + Returns a list of `FileEncoding` tuples with the detected encodings ordered + by confidence. + + Args: + file_path: The path to the file to detect the encoding for. + timeout: The timeout in seconds for the encoding detection. + """ + import chardet + + def read_and_detect(file_path: str) -> list[dict]: + rawdata = Path(file_path).read_bytes() + return cast(list[dict], chardet.detect_all(rawdata)) + + with concurrent.futures.ThreadPoolExecutor() as executor: + future = executor.submit(read_and_detect, file_path) + try: + encodings = future.result(timeout=timeout) + except concurrent.futures.TimeoutError: + raise TimeoutError(f"Timeout reached while detecting encoding for {file_path}") + + if all(encoding["encoding"] is None for encoding in encodings): + raise RuntimeError(f"Could not detect encoding for {file_path}") + return [FileEncoding(**enc) for enc in encodings if enc["encoding"] is not None] diff --git a/api/core/rag/extractor/html_extractor.py b/api/core/rag/extractor/html_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..350b522347b09dbb5bbc7fde136b05304c3e68e0 --- /dev/null +++ b/api/core/rag/extractor/html_extractor.py @@ -0,0 +1,32 @@ +"""Abstract interface for document loader implementations.""" + +from bs4 import BeautifulSoup # type: ignore + +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.models.document import Document + + +class HtmlExtractor(BaseExtractor): + """ + Load html files. + + + Args: + file_path: Path to the file to load. + """ + + def __init__(self, file_path: str): + """Initialize with file path.""" + self._file_path = file_path + + def extract(self) -> list[Document]: + return [Document(page_content=self._load_as_text())] + + def _load_as_text(self) -> str: + text: str = "" + with open(self._file_path, "rb") as fp: + soup = BeautifulSoup(fp, "html.parser") + text = soup.get_text() + text = text.strip() if text else "" + + return text diff --git a/api/core/rag/extractor/jina_reader_extractor.py b/api/core/rag/extractor/jina_reader_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..5b780af126b309883d4770cf085427b07c20e9b4 --- /dev/null +++ b/api/core/rag/extractor/jina_reader_extractor.py @@ -0,0 +1,35 @@ +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.models.document import Document +from services.website_service import WebsiteService + + +class JinaReaderWebExtractor(BaseExtractor): + """ + Crawl and scrape websites and return content in clean llm-ready markdown. + """ + + def __init__(self, url: str, job_id: str, tenant_id: str, mode: str = "crawl", only_main_content: bool = False): + """Initialize with url, api_key, base_url and mode.""" + self._url = url + self.job_id = job_id + self.tenant_id = tenant_id + self.mode = mode + self.only_main_content = only_main_content + + def extract(self) -> list[Document]: + """Extract content from the URL.""" + documents = [] + if self.mode == "crawl": + crawl_data = WebsiteService.get_crawl_url_data(self.job_id, "jinareader", self._url, self.tenant_id) + if crawl_data is None: + return [] + document = Document( + page_content=crawl_data.get("content", ""), + metadata={ + "source_url": crawl_data.get("url"), + "description": crawl_data.get("description"), + "title": crawl_data.get("title"), + }, + ) + documents.append(document) + return documents diff --git a/api/core/rag/extractor/markdown_extractor.py b/api/core/rag/extractor/markdown_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..849852ac23819ac9c3b3354b78be6c2a209ab42e --- /dev/null +++ b/api/core/rag/extractor/markdown_extractor.py @@ -0,0 +1,127 @@ +"""Abstract interface for document loader implementations.""" + +import re +from pathlib import Path +from typing import Optional, cast + +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.extractor.helpers import detect_file_encodings +from core.rag.models.document import Document + + +class MarkdownExtractor(BaseExtractor): + """Load Markdown files. + + + Args: + file_path: Path to the file to load. + """ + + def __init__( + self, + file_path: str, + remove_hyperlinks: bool = False, + remove_images: bool = False, + encoding: Optional[str] = None, + autodetect_encoding: bool = True, + ): + """Initialize with file path.""" + self._file_path = file_path + self._remove_hyperlinks = remove_hyperlinks + self._remove_images = remove_images + self._encoding = encoding + self._autodetect_encoding = autodetect_encoding + + def extract(self) -> list[Document]: + """Load from file path.""" + tups = self.parse_tups(self._file_path) + documents = [] + for header, value in tups: + value = value.strip() + if header is None: + documents.append(Document(page_content=value)) + else: + documents.append(Document(page_content=f"\n\n{header}\n{value}")) + + return documents + + def markdown_to_tups(self, markdown_text: str) -> list[tuple[Optional[str], str]]: + """Convert a markdown file to a dictionary. + + The keys are the headers and the values are the text under each header. + + """ + markdown_tups: list[tuple[Optional[str], str]] = [] + lines = markdown_text.split("\n") + + current_header = None + current_text = "" + code_block_flag = False + + for line in lines: + if line.startswith("```"): + code_block_flag = not code_block_flag + current_text += line + "\n" + continue + if code_block_flag: + current_text += line + "\n" + continue + header_match = re.match(r"^#+\s", line) + if header_match: + if current_header is not None: + markdown_tups.append((current_header, current_text)) + + current_header = line + current_text = "" + else: + current_text += line + "\n" + markdown_tups.append((current_header, current_text)) + + if current_header is not None: + # pass linting, assert keys are defined + markdown_tups = [ + (re.sub(r"#", "", cast(str, key)).strip(), re.sub(r"<.*?>", "", value)) for key, value in markdown_tups + ] + else: + markdown_tups = [(key, re.sub("\n", "", value)) for key, value in markdown_tups] + + return markdown_tups + + def remove_images(self, content: str) -> str: + """Get a dictionary of a markdown file from its path.""" + pattern = r"!{1}\[\[(.*)\]\]" + content = re.sub(pattern, "", content) + return content + + def remove_hyperlinks(self, content: str) -> str: + """Get a dictionary of a markdown file from its path.""" + pattern = r"\[(.*?)\]\((.*?)\)" + content = re.sub(pattern, r"\1", content) + return content + + def parse_tups(self, filepath: str) -> list[tuple[Optional[str], str]]: + """Parse file into tuples.""" + content = "" + try: + content = Path(filepath).read_text(encoding=self._encoding) + except UnicodeDecodeError as e: + if self._autodetect_encoding: + detected_encodings = detect_file_encodings(filepath) + for encoding in detected_encodings: + try: + content = Path(filepath).read_text(encoding=encoding.encoding) + break + except UnicodeDecodeError: + continue + else: + raise RuntimeError(f"Error loading {filepath}") from e + except Exception as e: + raise RuntimeError(f"Error loading {filepath}") from e + + if self._remove_hyperlinks: + content = self.remove_hyperlinks(content) + + if self._remove_images: + content = self.remove_images(content) + + return self.markdown_to_tups(content) diff --git a/api/core/rag/extractor/notion_extractor.py b/api/core/rag/extractor/notion_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..7ab248199a6a01870cf0b89e47d26a59a699f9a3 --- /dev/null +++ b/api/core/rag/extractor/notion_extractor.py @@ -0,0 +1,364 @@ +import json +import logging +from typing import Any, Optional, cast + +import requests + +from configs import dify_config +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.models.document import Document +from extensions.ext_database import db +from models.dataset import Document as DocumentModel +from models.source import DataSourceOauthBinding + +logger = logging.getLogger(__name__) + +BLOCK_CHILD_URL_TMPL = "https://api.notion.com/v1/blocks/{block_id}/children" +DATABASE_URL_TMPL = "https://api.notion.com/v1/databases/{database_id}/query" +SEARCH_URL = "https://api.notion.com/v1/search" + +RETRIEVE_PAGE_URL_TMPL = "https://api.notion.com/v1/pages/{page_id}" +RETRIEVE_DATABASE_URL_TMPL = "https://api.notion.com/v1/databases/{database_id}" +# if user want split by headings, use the corresponding splitter +HEADING_SPLITTER = { + "heading_1": "# ", + "heading_2": "## ", + "heading_3": "### ", +} + + +class NotionExtractor(BaseExtractor): + def __init__( + self, + notion_workspace_id: str, + notion_obj_id: str, + notion_page_type: str, + tenant_id: str, + document_model: Optional[DocumentModel] = None, + notion_access_token: Optional[str] = None, + ): + self._notion_access_token = None + self._document_model = document_model + self._notion_workspace_id = notion_workspace_id + self._notion_obj_id = notion_obj_id + self._notion_page_type = notion_page_type + if notion_access_token: + self._notion_access_token = notion_access_token + else: + self._notion_access_token = self._get_access_token(tenant_id, self._notion_workspace_id) + if not self._notion_access_token: + integration_token = dify_config.NOTION_INTEGRATION_TOKEN + if integration_token is None: + raise ValueError( + "Must specify `integration_token` or set environment variable `NOTION_INTEGRATION_TOKEN`." + ) + + self._notion_access_token = integration_token + + def extract(self) -> list[Document]: + self.update_last_edited_time(self._document_model) + + text_docs = self._load_data_as_documents(self._notion_obj_id, self._notion_page_type) + + return text_docs + + def _load_data_as_documents(self, notion_obj_id: str, notion_page_type: str) -> list[Document]: + docs = [] + if notion_page_type == "database": + # get all the pages in the database + page_text_documents = self._get_notion_database_data(notion_obj_id) + docs.extend(page_text_documents) + elif notion_page_type == "page": + page_text_list = self._get_notion_block_data(notion_obj_id) + docs.append(Document(page_content="\n".join(page_text_list))) + else: + raise ValueError("notion page type not supported") + + return docs + + def _get_notion_database_data(self, database_id: str, query_dict: dict[str, Any] = {}) -> list[Document]: + """Get all the pages from a Notion database.""" + assert self._notion_access_token is not None, "Notion access token is required" + res = requests.post( + DATABASE_URL_TMPL.format(database_id=database_id), + headers={ + "Authorization": "Bearer " + self._notion_access_token, + "Content-Type": "application/json", + "Notion-Version": "2022-06-28", + }, + json=query_dict, + ) + + data = res.json() + + database_content = [] + if "results" not in data or data["results"] is None: + return [] + for result in data["results"]: + properties = result["properties"] + data = {} + value: Any + for property_name, property_value in properties.items(): + type = property_value["type"] + if type == "multi_select": + value = [] + multi_select_list = property_value[type] + for multi_select in multi_select_list: + value.append(multi_select["name"]) + elif type in {"rich_text", "title"}: + if len(property_value[type]) > 0: + value = property_value[type][0]["plain_text"] + else: + value = "" + elif type in {"select", "status"}: + if property_value[type]: + value = property_value[type]["name"] + else: + value = "" + else: + value = property_value[type] + data[property_name] = value + row_dict = {k: v for k, v in data.items() if v} + row_content = "" + for key, value in row_dict.items(): + if isinstance(value, dict): + value_dict = {k: v for k, v in value.items() if v} + value_content = "".join(f"{k}:{v} " for k, v in value_dict.items()) + row_content = row_content + f"{key}:{value_content}\n" + else: + row_content = row_content + f"{key}:{value}\n" + database_content.append(row_content) + + return [Document(page_content="\n".join(database_content))] + + def _get_notion_block_data(self, page_id: str) -> list[str]: + assert self._notion_access_token is not None, "Notion access token is required" + result_lines_arr = [] + start_cursor = None + block_url = BLOCK_CHILD_URL_TMPL.format(block_id=page_id) + while True: + query_dict: dict[str, Any] = {} if not start_cursor else {"start_cursor": start_cursor} + try: + res = requests.request( + "GET", + block_url, + headers={ + "Authorization": "Bearer " + self._notion_access_token, + "Content-Type": "application/json", + "Notion-Version": "2022-06-28", + }, + params=query_dict, + ) + if res.status_code != 200: + raise ValueError(f"Error fetching Notion block data: {res.text}") + data = res.json() + except requests.RequestException as e: + raise ValueError("Error fetching Notion block data") from e + if "results" not in data or not isinstance(data["results"], list): + raise ValueError("Error fetching Notion block data") + for result in data["results"]: + result_type = result["type"] + result_obj = result[result_type] + cur_result_text_arr = [] + if result_type == "table": + result_block_id = result["id"] + text = self._read_table_rows(result_block_id) + text += "\n\n" + result_lines_arr.append(text) + else: + if "rich_text" in result_obj: + for rich_text in result_obj["rich_text"]: + # skip if doesn't have text object + if "text" in rich_text: + text = rich_text["text"]["content"] + cur_result_text_arr.append(text) + + result_block_id = result["id"] + has_children = result["has_children"] + block_type = result["type"] + if has_children and block_type != "child_page": + children_text = self._read_block(result_block_id, num_tabs=1) + cur_result_text_arr.append(children_text) + + cur_result_text = "\n".join(cur_result_text_arr) + if result_type in HEADING_SPLITTER: + result_lines_arr.append(f"{HEADING_SPLITTER[result_type]}{cur_result_text}") + else: + result_lines_arr.append(cur_result_text + "\n\n") + + if data["next_cursor"] is None: + break + else: + start_cursor = data["next_cursor"] + return result_lines_arr + + def _read_block(self, block_id: str, num_tabs: int = 0) -> str: + """Read a block.""" + assert self._notion_access_token is not None, "Notion access token is required" + result_lines_arr = [] + start_cursor = None + block_url = BLOCK_CHILD_URL_TMPL.format(block_id=block_id) + while True: + query_dict: dict[str, Any] = {} if not start_cursor else {"start_cursor": start_cursor} + + res = requests.request( + "GET", + block_url, + headers={ + "Authorization": "Bearer " + self._notion_access_token, + "Content-Type": "application/json", + "Notion-Version": "2022-06-28", + }, + params=query_dict, + ) + data = res.json() + if "results" not in data or data["results"] is None: + break + for result in data["results"]: + result_type = result["type"] + result_obj = result[result_type] + cur_result_text_arr = [] + if result_type == "table": + result_block_id = result["id"] + text = self._read_table_rows(result_block_id) + result_lines_arr.append(text) + else: + if "rich_text" in result_obj: + for rich_text in result_obj["rich_text"]: + # skip if doesn't have text object + if "text" in rich_text: + text = rich_text["text"]["content"] + prefix = "\t" * num_tabs + cur_result_text_arr.append(prefix + text) + result_block_id = result["id"] + has_children = result["has_children"] + block_type = result["type"] + if has_children and block_type != "child_page": + children_text = self._read_block(result_block_id, num_tabs=num_tabs + 1) + cur_result_text_arr.append(children_text) + + cur_result_text = "\n".join(cur_result_text_arr) + if result_type in HEADING_SPLITTER: + result_lines_arr.append(f"{HEADING_SPLITTER[result_type]}{cur_result_text}") + else: + result_lines_arr.append(cur_result_text + "\n\n") + + if data["next_cursor"] is None: + break + else: + start_cursor = data["next_cursor"] + + result_lines = "\n".join(result_lines_arr) + return result_lines + + def _read_table_rows(self, block_id: str) -> str: + """Read table rows.""" + assert self._notion_access_token is not None, "Notion access token is required" + done = False + result_lines_arr = [] + start_cursor = None + block_url = BLOCK_CHILD_URL_TMPL.format(block_id=block_id) + while not done: + query_dict: dict[str, Any] = {} if not start_cursor else {"start_cursor": start_cursor} + + res = requests.request( + "GET", + block_url, + headers={ + "Authorization": "Bearer " + self._notion_access_token, + "Content-Type": "application/json", + "Notion-Version": "2022-06-28", + }, + params=query_dict, + ) + data = res.json() + # get table headers text + table_header_cell_texts = [] + table_header_cells = data["results"][0]["table_row"]["cells"] + for table_header_cell in table_header_cells: + if table_header_cell: + for table_header_cell_text in table_header_cell: + text = table_header_cell_text["text"]["content"] + table_header_cell_texts.append(text) + else: + table_header_cell_texts.append("") + # Initialize Markdown table with headers + markdown_table = "| " + " | ".join(table_header_cell_texts) + " |\n" + markdown_table += "| " + " | ".join(["---"] * len(table_header_cell_texts)) + " |\n" + + # Process data to format each row in Markdown table format + results = data["results"] + for i in range(len(results) - 1): + column_texts = [] + table_column_cells = data["results"][i + 1]["table_row"]["cells"] + for j in range(len(table_column_cells)): + if table_column_cells[j]: + for table_column_cell_text in table_column_cells[j]: + column_text = table_column_cell_text["text"]["content"] + column_texts.append(column_text) + # Add row to Markdown table + markdown_table += "| " + " | ".join(column_texts) + " |\n" + result_lines_arr.append(markdown_table) + if data["next_cursor"] is None: + done = True + break + else: + start_cursor = data["next_cursor"] + + result_lines = "\n".join(result_lines_arr) + return result_lines + + def update_last_edited_time(self, document_model: Optional[DocumentModel]): + if not document_model: + return + + last_edited_time = self.get_notion_last_edited_time() + data_source_info = document_model.data_source_info_dict + data_source_info["last_edited_time"] = last_edited_time + update_params = {DocumentModel.data_source_info: json.dumps(data_source_info)} + + DocumentModel.query.filter_by(id=document_model.id).update(update_params) + db.session.commit() + + def get_notion_last_edited_time(self) -> str: + assert self._notion_access_token is not None, "Notion access token is required" + obj_id = self._notion_obj_id + page_type = self._notion_page_type + if page_type == "database": + retrieve_page_url = RETRIEVE_DATABASE_URL_TMPL.format(database_id=obj_id) + else: + retrieve_page_url = RETRIEVE_PAGE_URL_TMPL.format(page_id=obj_id) + + query_dict: dict[str, Any] = {} + + res = requests.request( + "GET", + retrieve_page_url, + headers={ + "Authorization": "Bearer " + self._notion_access_token, + "Content-Type": "application/json", + "Notion-Version": "2022-06-28", + }, + json=query_dict, + ) + + data = res.json() + return cast(str, data["last_edited_time"]) + + @classmethod + def _get_access_token(cls, tenant_id: str, notion_workspace_id: str) -> str: + data_source_binding = DataSourceOauthBinding.query.filter( + db.and_( + DataSourceOauthBinding.tenant_id == tenant_id, + DataSourceOauthBinding.provider == "notion", + DataSourceOauthBinding.disabled == False, + DataSourceOauthBinding.source_info["workspace_id"] == f'"{notion_workspace_id}"', + ) + ).first() + + if not data_source_binding: + raise Exception( + f"No notion data source binding found for tenant {tenant_id} and notion workspace {notion_workspace_id}" + ) + + return cast(str, data_source_binding.access_token) diff --git a/api/core/rag/extractor/pdf_extractor.py b/api/core/rag/extractor/pdf_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..04033dec3fa56de6fedf2cb07b384b00015a9561 --- /dev/null +++ b/api/core/rag/extractor/pdf_extractor.py @@ -0,0 +1,68 @@ +"""Abstract interface for document loader implementations.""" + +from collections.abc import Iterator +from typing import Optional, cast + +from core.rag.extractor.blob.blob import Blob +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.models.document import Document +from extensions.ext_storage import storage + + +class PdfExtractor(BaseExtractor): + """Load pdf files. + + + Args: + file_path: Path to the file to load. + """ + + def __init__(self, file_path: str, file_cache_key: Optional[str] = None): + """Initialize with file path.""" + self._file_path = file_path + self._file_cache_key = file_cache_key + + def extract(self) -> list[Document]: + plaintext_file_exists = False + if self._file_cache_key: + try: + text = cast(bytes, storage.load(self._file_cache_key)).decode("utf-8") + plaintext_file_exists = True + return [Document(page_content=text)] + except FileNotFoundError: + pass + documents = list(self.load()) + text_list = [] + for document in documents: + text_list.append(document.page_content) + text = "\n\n".join(text_list) + + # save plaintext file for caching + if not plaintext_file_exists and self._file_cache_key: + storage.save(self._file_cache_key, text.encode("utf-8")) + + return documents + + def load( + self, + ) -> Iterator[Document]: + """Lazy load given path as pages.""" + blob = Blob.from_path(self._file_path) + yield from self.parse(blob) + + def parse(self, blob: Blob) -> Iterator[Document]: + """Lazily parse the blob.""" + import pypdfium2 # type: ignore + + with blob.as_bytes_io() as file_path: + pdf_reader = pypdfium2.PdfDocument(file_path, autoclose=True) + try: + for page_number, page in enumerate(pdf_reader): + text_page = page.get_textpage() + content = text_page.get_text_range() + text_page.close() + page.close() + metadata = {"source": blob.source, "page": page_number} + yield Document(page_content=content, metadata=metadata) + finally: + pdf_reader.close() diff --git a/api/core/rag/extractor/text_extractor.py b/api/core/rag/extractor/text_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..b2b51d71d73a16a2d293b163909fc2c99b84d2f2 --- /dev/null +++ b/api/core/rag/extractor/text_extractor.py @@ -0,0 +1,45 @@ +"""Abstract interface for document loader implementations.""" + +from pathlib import Path +from typing import Optional + +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.extractor.helpers import detect_file_encodings +from core.rag.models.document import Document + + +class TextExtractor(BaseExtractor): + """Load text files. + + + Args: + file_path: Path to the file to load. + """ + + def __init__(self, file_path: str, encoding: Optional[str] = None, autodetect_encoding: bool = False): + """Initialize with file path.""" + self._file_path = file_path + self._encoding = encoding + self._autodetect_encoding = autodetect_encoding + + def extract(self) -> list[Document]: + """Load from file path.""" + text = "" + try: + text = Path(self._file_path).read_text(encoding=self._encoding) + except UnicodeDecodeError as e: + if self._autodetect_encoding: + detected_encodings = detect_file_encodings(self._file_path) + for encoding in detected_encodings: + try: + text = Path(self._file_path).read_text(encoding=encoding.encoding) + break + except UnicodeDecodeError: + continue + else: + raise RuntimeError(f"Error loading {self._file_path}") from e + except Exception as e: + raise RuntimeError(f"Error loading {self._file_path}") from e + + metadata = {"source": self._file_path} + return [Document(page_content=text, metadata=metadata)] diff --git a/api/core/rag/extractor/unstructured/unstructured_doc_extractor.py b/api/core/rag/extractor/unstructured/unstructured_doc_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..a525c9e9e3c443875bf25c015a065aae5eeb46e4 --- /dev/null +++ b/api/core/rag/extractor/unstructured/unstructured_doc_extractor.py @@ -0,0 +1,59 @@ +import logging +import os + +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.models.document import Document + +logger = logging.getLogger(__name__) + + +class UnstructuredWordExtractor(BaseExtractor): + """Loader that uses unstructured to load word documents.""" + + def __init__( + self, + file_path: str, + api_url: str, + ): + """Initialize with file path.""" + self._file_path = file_path + self._api_url = api_url + + def extract(self) -> list[Document]: + from unstructured.__version__ import __version__ as __unstructured_version__ + from unstructured.file_utils.filetype import FileType, detect_filetype + + unstructured_version = tuple(int(x) for x in __unstructured_version__.split(".")) + # check the file extension + try: + import magic # noqa: F401 + + is_doc = detect_filetype(self._file_path) == FileType.DOC + except ImportError: + _, extension = os.path.splitext(str(self._file_path)) + is_doc = extension == ".doc" + + if is_doc and unstructured_version < (0, 4, 11): + raise ValueError( + f"You are on unstructured version {__unstructured_version__}. " + "Partitioning .doc files is only supported in unstructured>=0.4.11. " + "Please upgrade the unstructured package and try again." + ) + + if is_doc: + from unstructured.partition.doc import partition_doc + + elements = partition_doc(filename=self._file_path) + else: + from unstructured.partition.docx import partition_docx + + elements = partition_docx(filename=self._file_path) + + from unstructured.chunking.title import chunk_by_title + + chunks = chunk_by_title(elements, max_characters=2000, combine_text_under_n_chars=2000) + documents = [] + for chunk in chunks: + text = chunk.text.strip() + documents.append(Document(page_content=text)) + return documents diff --git a/api/core/rag/extractor/unstructured/unstructured_eml_extractor.py b/api/core/rag/extractor/unstructured/unstructured_eml_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..f1fa5dde5c19a6864e603a7d1e47135cd78204f7 --- /dev/null +++ b/api/core/rag/extractor/unstructured/unstructured_eml_extractor.py @@ -0,0 +1,56 @@ +import base64 +import logging +from typing import Optional + +from bs4 import BeautifulSoup # type: ignore + +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.models.document import Document + +logger = logging.getLogger(__name__) + + +class UnstructuredEmailExtractor(BaseExtractor): + """Load eml files. + Args: + file_path: Path to the file to load. + """ + + def __init__(self, file_path: str, api_url: Optional[str] = None, api_key: str = ""): + """Initialize with file path.""" + self._file_path = file_path + self._api_url = api_url + self._api_key = api_key + + def extract(self) -> list[Document]: + if self._api_url: + from unstructured.partition.api import partition_via_api + + elements = partition_via_api(filename=self._file_path, api_url=self._api_url, api_key=self._api_key) + else: + from unstructured.partition.email import partition_email + + elements = partition_email(filename=self._file_path) + + # noinspection PyBroadException + try: + for element in elements: + element_text = element.text.strip() + + padding_needed = 4 - len(element_text) % 4 + element_text += "=" * padding_needed + + element_decode = base64.b64decode(element_text) + soup = BeautifulSoup(element_decode.decode("utf-8"), "html.parser") + element.text = soup.get_text() + except Exception: + pass + + from unstructured.chunking.title import chunk_by_title + + chunks = chunk_by_title(elements, max_characters=2000, combine_text_under_n_chars=2000) + documents = [] + for chunk in chunks: + text = chunk.text.strip() + documents.append(Document(page_content=text)) + return documents diff --git a/api/core/rag/extractor/unstructured/unstructured_epub_extractor.py b/api/core/rag/extractor/unstructured/unstructured_epub_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..35ca686f62dbf8d90b4756ec9937e42feb75cdca --- /dev/null +++ b/api/core/rag/extractor/unstructured/unstructured_epub_extractor.py @@ -0,0 +1,47 @@ +import logging +from typing import Optional + +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.models.document import Document + +logger = logging.getLogger(__name__) + + +class UnstructuredEpubExtractor(BaseExtractor): + """Load epub files. + + + Args: + file_path: Path to the file to load. + """ + + def __init__( + self, + file_path: str, + api_url: Optional[str] = None, + api_key: str = "", + ): + """Initialize with file path.""" + self._file_path = file_path + self._api_url = api_url + self._api_key = api_key + + def extract(self) -> list[Document]: + if self._api_url: + from unstructured.partition.api import partition_via_api + + elements = partition_via_api(filename=self._file_path, api_url=self._api_url, api_key=self._api_key) + else: + from unstructured.partition.epub import partition_epub + + elements = partition_epub(filename=self._file_path, xml_keep_tags=True) + + from unstructured.chunking.title import chunk_by_title + + chunks = chunk_by_title(elements, max_characters=2000, combine_text_under_n_chars=2000) + documents = [] + for chunk in chunks: + text = chunk.text.strip() + documents.append(Document(page_content=text)) + + return documents diff --git a/api/core/rag/extractor/unstructured/unstructured_markdown_extractor.py b/api/core/rag/extractor/unstructured/unstructured_markdown_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..d5418e612ab594e0c5a0c5fafac09b1f2b61ab23 --- /dev/null +++ b/api/core/rag/extractor/unstructured/unstructured_markdown_extractor.py @@ -0,0 +1,51 @@ +import logging +from typing import Optional + +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.models.document import Document + +logger = logging.getLogger(__name__) + + +class UnstructuredMarkdownExtractor(BaseExtractor): + """Load md files. + + + Args: + file_path: Path to the file to load. + + remove_hyperlinks: Whether to remove hyperlinks from the text. + + remove_images: Whether to remove images from the text. + + encoding: File encoding to use. If `None`, the file will be loaded + with the default system encoding. + + autodetect_encoding: Whether to try to autodetect the file encoding + if the specified encoding fails. + """ + + def __init__(self, file_path: str, api_url: Optional[str] = None, api_key: str = ""): + """Initialize with file path.""" + self._file_path = file_path + self._api_url = api_url + self._api_key = api_key + + def extract(self) -> list[Document]: + if self._api_url: + from unstructured.partition.api import partition_via_api + + elements = partition_via_api(filename=self._file_path, api_url=self._api_url, api_key=self._api_key) + else: + from unstructured.partition.md import partition_md + + elements = partition_md(filename=self._file_path) + from unstructured.chunking.title import chunk_by_title + + chunks = chunk_by_title(elements, max_characters=2000, combine_text_under_n_chars=2000) + documents = [] + for chunk in chunks: + text = chunk.text.strip() + documents.append(Document(page_content=text)) + + return documents diff --git a/api/core/rag/extractor/unstructured/unstructured_msg_extractor.py b/api/core/rag/extractor/unstructured/unstructured_msg_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..d363449c29dad5821d83bac2f6f28773e82ca4a5 --- /dev/null +++ b/api/core/rag/extractor/unstructured/unstructured_msg_extractor.py @@ -0,0 +1,41 @@ +import logging +from typing import Optional + +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.models.document import Document + +logger = logging.getLogger(__name__) + + +class UnstructuredMsgExtractor(BaseExtractor): + """Load msg files. + + + Args: + file_path: Path to the file to load. + """ + + def __init__(self, file_path: str, api_url: Optional[str] = None, api_key: str = ""): + """Initialize with file path.""" + self._file_path = file_path + self._api_url = api_url + self._api_key = api_key + + def extract(self) -> list[Document]: + if self._api_url: + from unstructured.partition.api import partition_via_api + + elements = partition_via_api(filename=self._file_path, api_url=self._api_url, api_key=self._api_key) + else: + from unstructured.partition.msg import partition_msg + + elements = partition_msg(filename=self._file_path) + from unstructured.chunking.title import chunk_by_title + + chunks = chunk_by_title(elements, max_characters=2000, combine_text_under_n_chars=2000) + documents = [] + for chunk in chunks: + text = chunk.text.strip() + documents.append(Document(page_content=text)) + + return documents diff --git a/api/core/rag/extractor/unstructured/unstructured_pdf_extractor.py b/api/core/rag/extractor/unstructured/unstructured_pdf_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..dd8a979e70998921b2e36b5edb3d351db698bedc --- /dev/null +++ b/api/core/rag/extractor/unstructured/unstructured_pdf_extractor.py @@ -0,0 +1,47 @@ +import logging + +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.models.document import Document + +logger = logging.getLogger(__name__) + + +class UnstructuredPDFExtractor(BaseExtractor): + """Load pdf files. + + + Args: + file_path: Path to the file to load. + + api_url: Unstructured API URL + + api_key: Unstructured API Key + """ + + def __init__(self, file_path: str, api_url: str, api_key: str): + """Initialize with file path.""" + self._file_path = file_path + self._api_url = api_url + self._api_key = api_key + + def extract(self) -> list[Document]: + if self._api_url: + from unstructured.partition.api import partition_via_api + + elements = partition_via_api( + filename=self._file_path, api_url=self._api_url, api_key=self._api_key, strategy="auto" + ) + else: + from unstructured.partition.pdf import partition_pdf + + elements = partition_pdf(filename=self._file_path, strategy="auto") + + from unstructured.chunking.title import chunk_by_title + + chunks = chunk_by_title(elements, max_characters=2000, combine_text_under_n_chars=2000) + documents = [] + for chunk in chunks: + text = chunk.text.strip() + documents.append(Document(page_content=text)) + + return documents diff --git a/api/core/rag/extractor/unstructured/unstructured_ppt_extractor.py b/api/core/rag/extractor/unstructured/unstructured_ppt_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..ecc272a2f00f0fc8860d45d04fb9cc3c55a80e25 --- /dev/null +++ b/api/core/rag/extractor/unstructured/unstructured_ppt_extractor.py @@ -0,0 +1,47 @@ +import logging +from typing import Optional + +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.models.document import Document + +logger = logging.getLogger(__name__) + + +class UnstructuredPPTExtractor(BaseExtractor): + """Load ppt files. + + + Args: + file_path: Path to the file to load. + """ + + def __init__(self, file_path: str, api_url: Optional[str] = None, api_key: str = ""): + """Initialize with file path.""" + self._file_path = file_path + self._api_url = api_url + self._api_key = api_key + + def extract(self) -> list[Document]: + if self._api_url: + from unstructured.partition.api import partition_via_api + + elements = partition_via_api(filename=self._file_path, api_url=self._api_url, api_key=self._api_key) + else: + raise NotImplementedError("Unstructured API Url is not configured") + text_by_page: dict[int, str] = {} + for element in elements: + page = element.metadata.page_number + if page is None: + continue + text = element.text + if page in text_by_page: + text_by_page[page] += "\n" + text + else: + text_by_page[page] = text + + combined_texts = list(text_by_page.values()) + documents = [] + for combined_text in combined_texts: + text = combined_text.strip() + documents.append(Document(page_content=text)) + return documents diff --git a/api/core/rag/extractor/unstructured/unstructured_pptx_extractor.py b/api/core/rag/extractor/unstructured/unstructured_pptx_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..e7bf6fd2e6ea8809630815f4f673b4ed993b7d43 --- /dev/null +++ b/api/core/rag/extractor/unstructured/unstructured_pptx_extractor.py @@ -0,0 +1,49 @@ +import logging +from typing import Optional + +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.models.document import Document + +logger = logging.getLogger(__name__) + + +class UnstructuredPPTXExtractor(BaseExtractor): + """Load pptx files. + + + Args: + file_path: Path to the file to load. + """ + + def __init__(self, file_path: str, api_url: Optional[str] = None, api_key: str = ""): + """Initialize with file path.""" + self._file_path = file_path + self._api_url = api_url + self._api_key = api_key + + def extract(self) -> list[Document]: + if self._api_url: + from unstructured.partition.api import partition_via_api + + elements = partition_via_api(filename=self._file_path, api_url=self._api_url, api_key=self._api_key) + else: + from unstructured.partition.pptx import partition_pptx + + elements = partition_pptx(filename=self._file_path) + text_by_page: dict[int, str] = {} + for element in elements: + page = element.metadata.page_number + text = element.text + if page is not None: + if page in text_by_page: + text_by_page[page] += "\n" + text + else: + text_by_page[page] = text + + combined_texts = list(text_by_page.values()) + documents = [] + for combined_text in combined_texts: + text = combined_text.strip() + documents.append(Document(page_content=text)) + + return documents diff --git a/api/core/rag/extractor/unstructured/unstructured_text_extractor.py b/api/core/rag/extractor/unstructured/unstructured_text_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..22dfdd20752cbfc8d9975f8c0e38ad2493770c8d --- /dev/null +++ b/api/core/rag/extractor/unstructured/unstructured_text_extractor.py @@ -0,0 +1,34 @@ +import logging + +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.models.document import Document + +logger = logging.getLogger(__name__) + + +class UnstructuredTextExtractor(BaseExtractor): + """Load msg files. + + + Args: + file_path: Path to the file to load. + """ + + def __init__(self, file_path: str, api_url: str): + """Initialize with file path.""" + self._file_path = file_path + self._api_url = api_url + + def extract(self) -> list[Document]: + from unstructured.partition.text import partition_text + + elements = partition_text(filename=self._file_path) + from unstructured.chunking.title import chunk_by_title + + chunks = chunk_by_title(elements, max_characters=2000, combine_text_under_n_chars=2000) + documents = [] + for chunk in chunks: + text = chunk.text.strip() + documents.append(Document(page_content=text)) + + return documents diff --git a/api/core/rag/extractor/unstructured/unstructured_xml_extractor.py b/api/core/rag/extractor/unstructured/unstructured_xml_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..916cdc3f2ba0f8fcd9a7be47490680d7e6db6a07 --- /dev/null +++ b/api/core/rag/extractor/unstructured/unstructured_xml_extractor.py @@ -0,0 +1,42 @@ +import logging +from typing import Optional + +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.models.document import Document + +logger = logging.getLogger(__name__) + + +class UnstructuredXmlExtractor(BaseExtractor): + """Load xml files. + + + Args: + file_path: Path to the file to load. + """ + + def __init__(self, file_path: str, api_url: Optional[str] = None, api_key: str = ""): + """Initialize with file path.""" + self._file_path = file_path + self._api_url = api_url + self._api_key = api_key + + def extract(self) -> list[Document]: + if self._api_url: + from unstructured.partition.api import partition_via_api + + elements = partition_via_api(filename=self._file_path, api_url=self._api_url, api_key=self._api_key) + else: + from unstructured.partition.xml import partition_xml + + elements = partition_xml(filename=self._file_path, xml_keep_tags=True) + + from unstructured.chunking.title import chunk_by_title + + chunks = chunk_by_title(elements, max_characters=2000, combine_text_under_n_chars=2000) + documents = [] + for chunk in chunks: + text = chunk.text.strip() + documents.append(Document(page_content=text)) + + return documents diff --git a/api/core/rag/extractor/word_extractor.py b/api/core/rag/extractor/word_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..d93de5fef948d491e54d7569668f13c7d58fefac --- /dev/null +++ b/api/core/rag/extractor/word_extractor.py @@ -0,0 +1,277 @@ +"""Abstract interface for document loader implementations.""" + +import datetime +import logging +import mimetypes +import os +import re +import tempfile +import uuid +from urllib.parse import urlparse +from xml.etree import ElementTree + +import requests +from docx import Document as DocxDocument + +from configs import dify_config +from core.helper import ssrf_proxy +from core.rag.extractor.extractor_base import BaseExtractor +from core.rag.models.document import Document +from extensions.ext_database import db +from extensions.ext_storage import storage +from models.enums import CreatedByRole +from models.model import UploadFile + +logger = logging.getLogger(__name__) + + +class WordExtractor(BaseExtractor): + """Load docx files. + + Args: + file_path: Path to the file to load. + """ + + def __init__(self, file_path: str, tenant_id: str, user_id: str): + """Initialize with file path.""" + self.file_path = file_path + self.tenant_id = tenant_id + self.user_id = user_id + + if "~" in self.file_path: + self.file_path = os.path.expanduser(self.file_path) + + # If the file is a web path, download it to a temporary file, and use that + if not os.path.isfile(self.file_path) and self._is_valid_url(self.file_path): + r = requests.get(self.file_path) + + if r.status_code != 200: + raise ValueError(f"Check the url of your file; returned status code {r.status_code}") + + self.web_path = self.file_path + # TODO: use a better way to handle the file + self.temp_file = tempfile.NamedTemporaryFile() # noqa: SIM115 + self.temp_file.write(r.content) + self.file_path = self.temp_file.name + elif not os.path.isfile(self.file_path): + raise ValueError(f"File path {self.file_path} is not a valid file or url") + + def __del__(self) -> None: + if hasattr(self, "temp_file"): + self.temp_file.close() + + def extract(self) -> list[Document]: + """Load given path as single page.""" + content = self.parse_docx(self.file_path, "storage") + return [ + Document( + page_content=content, + metadata={"source": self.file_path}, + ) + ] + + @staticmethod + def _is_valid_url(url: str) -> bool: + """Check if the url is valid.""" + parsed = urlparse(url) + return bool(parsed.netloc) and bool(parsed.scheme) + + def _extract_images_from_docx(self, doc, image_folder): + os.makedirs(image_folder, exist_ok=True) + image_count = 0 + image_map = {} + + for rel in doc.part.rels.values(): + if "image" in rel.target_ref: + image_count += 1 + if rel.is_external: + url = rel.reltype + response = ssrf_proxy.get(url) + if response.status_code == 200: + image_ext = mimetypes.guess_extension(response.headers["Content-Type"]) + if image_ext is None: + continue + file_uuid = str(uuid.uuid4()) + file_key = "image_files/" + self.tenant_id + "/" + file_uuid + "." + image_ext + mime_type, _ = mimetypes.guess_type(file_key) + storage.save(file_key, response.content) + else: + continue + else: + image_ext = rel.target_ref.split(".")[-1] + if image_ext is None: + continue + # user uuid as file name + file_uuid = str(uuid.uuid4()) + file_key = "image_files/" + self.tenant_id + "/" + file_uuid + "." + image_ext + mime_type, _ = mimetypes.guess_type(file_key) + + storage.save(file_key, rel.target_part.blob) + # save file to db + upload_file = UploadFile( + tenant_id=self.tenant_id, + storage_type=dify_config.STORAGE_TYPE, + key=file_key, + name=file_key, + size=0, + extension=str(image_ext), + mime_type=mime_type or "", + created_by=self.user_id, + created_by_role=CreatedByRole.ACCOUNT, + created_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None), + used=True, + used_by=self.user_id, + used_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None), + ) + + db.session.add(upload_file) + db.session.commit() + image_map[rel.target_part] = ( + f"![image]({dify_config.CONSOLE_API_URL}/files/{upload_file.id}/file-preview)" + ) + + return image_map + + def _table_to_markdown(self, table, image_map): + markdown = [] + # calculate the total number of columns + total_cols = max(len(row.cells) for row in table.rows) + + header_row = table.rows[0] + headers = self._parse_row(header_row, image_map, total_cols) + markdown.append("| " + " | ".join(headers) + " |") + markdown.append("| " + " | ".join(["---"] * total_cols) + " |") + + for row in table.rows[1:]: + row_cells = self._parse_row(row, image_map, total_cols) + markdown.append("| " + " | ".join(row_cells) + " |") + return "\n".join(markdown) + + def _parse_row(self, row, image_map, total_cols): + # Initialize a row, all of which are empty by default + row_cells = [""] * total_cols + col_index = 0 + for cell in row.cells: + # make sure the col_index is not out of range + while col_index < total_cols and row_cells[col_index] != "": + col_index += 1 + # if col_index is out of range the loop is jumped + if col_index >= total_cols: + break + cell_content = self._parse_cell(cell, image_map).strip() + cell_colspan = cell.grid_span or 1 + for i in range(cell_colspan): + if col_index + i < total_cols: + row_cells[col_index + i] = cell_content if i == 0 else "" + col_index += cell_colspan + return row_cells + + def _parse_cell(self, cell, image_map): + cell_content = [] + for paragraph in cell.paragraphs: + parsed_paragraph = self._parse_cell_paragraph(paragraph, image_map) + if parsed_paragraph: + cell_content.append(parsed_paragraph) + unique_content = list(dict.fromkeys(cell_content)) + return " ".join(unique_content) + + def _parse_cell_paragraph(self, paragraph, image_map): + paragraph_content = [] + for run in paragraph.runs: + if run.element.xpath(".//a:blip"): + for blip in run.element.xpath(".//a:blip"): + image_id = blip.get("{http://schemas.openxmlformats.org/officeDocument/2006/relationships}embed") + if not image_id: + continue + image_part = paragraph.part.rels[image_id].target_part + + if image_part in image_map: + image_link = image_map[image_part] + paragraph_content.append(image_link) + else: + paragraph_content.append(run.text) + return "".join(paragraph_content).strip() + + def _parse_paragraph(self, paragraph, image_map): + paragraph_content = [] + for run in paragraph.runs: + if run.element.xpath(".//a:blip"): + for blip in run.element.xpath(".//a:blip"): + embed_id = blip.get("{http://schemas.openxmlformats.org/officeDocument/2006/relationships}embed") + if embed_id: + rel_target = run.part.rels[embed_id].target_ref + if rel_target in image_map: + paragraph_content.append(image_map[rel_target]) + if run.text.strip(): + paragraph_content.append(run.text.strip()) + return " ".join(paragraph_content) if paragraph_content else "" + + def parse_docx(self, docx_path, image_folder): + doc = DocxDocument(docx_path) + os.makedirs(image_folder, exist_ok=True) + + content = [] + + image_map = self._extract_images_from_docx(doc, image_folder) + + hyperlinks_url = None + url_pattern = re.compile(r"http://[^\s+]+//|https://[^\s+]+") + for para in doc.paragraphs: + for run in para.runs: + if run.text and hyperlinks_url: + result = f" [{run.text}]({hyperlinks_url}) " + run.text = result + hyperlinks_url = None + if "HYPERLINK" in run.element.xml: + try: + xml = ElementTree.XML(run.element.xml) + x_child = [c for c in xml.iter() if c is not None] + for x in x_child: + if x_child is None: + continue + if x.tag.endswith("instrText"): + if x.text is None: + continue + for i in url_pattern.findall(x.text): + hyperlinks_url = str(i) + except Exception as e: + logger.exception("Failed to parse HYPERLINK xml") + + def parse_paragraph(paragraph): + paragraph_content = [] + for run in paragraph.runs: + if hasattr(run.element, "tag") and isinstance(run.element.tag, str) and run.element.tag.endswith("r"): + drawing_elements = run.element.findall( + ".//{http://schemas.openxmlformats.org/wordprocessingml/2006/main}drawing" + ) + for drawing in drawing_elements: + blip_elements = drawing.findall( + ".//{http://schemas.openxmlformats.org/drawingml/2006/main}blip" + ) + for blip in blip_elements: + embed_id = blip.get( + "{http://schemas.openxmlformats.org/officeDocument/2006/relationships}embed" + ) + if embed_id: + image_part = doc.part.related_parts.get(embed_id) + if image_part in image_map: + paragraph_content.append(image_map[image_part]) + if run.text.strip(): + paragraph_content.append(run.text.strip()) + return "".join(paragraph_content) if paragraph_content else "" + + paragraphs = doc.paragraphs.copy() + tables = doc.tables.copy() + for element in doc.element.body: + if hasattr(element, "tag"): + if isinstance(element.tag, str) and element.tag.endswith("p"): # paragraph + para = paragraphs.pop(0) + parsed_paragraph = parse_paragraph(para) + if parsed_paragraph.strip(): + content.append(parsed_paragraph) + else: + content.append("\n") + elif isinstance(element.tag, str) and element.tag.endswith("tbl"): # table + table = tables.pop(0) + content.append(self._table_to_markdown(table, image_map)) + return "\n".join(content) diff --git a/api/core/rag/index_processor/__init__.py b/api/core/rag/index_processor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/index_processor/constant/__init__.py b/api/core/rag/index_processor/constant/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/index_processor/constant/index_type.py b/api/core/rag/index_processor/constant/index_type.py new file mode 100644 index 0000000000000000000000000000000000000000..0845b58e25b5587cb1fae2d62f199d73bc6c18f6 --- /dev/null +++ b/api/core/rag/index_processor/constant/index_type.py @@ -0,0 +1,7 @@ +from enum import Enum + + +class IndexType(str, Enum): + PARAGRAPH_INDEX = "text_model" + QA_INDEX = "qa_model" + PARENT_CHILD_INDEX = "hierarchical_model" diff --git a/api/core/rag/index_processor/index_processor_base.py b/api/core/rag/index_processor/index_processor_base.py new file mode 100644 index 0000000000000000000000000000000000000000..2bcd1c79bb5dd44a6a16c9cde1bbd0613b7357f8 --- /dev/null +++ b/api/core/rag/index_processor/index_processor_base.py @@ -0,0 +1,84 @@ +"""Abstract interface for document loader implementations.""" + +from abc import ABC, abstractmethod +from typing import Optional + +from configs import dify_config +from core.model_manager import ModelInstance +from core.rag.extractor.entity.extract_setting import ExtractSetting +from core.rag.models.document import Document +from core.rag.splitter.fixed_text_splitter import ( + EnhanceRecursiveCharacterTextSplitter, + FixedRecursiveCharacterTextSplitter, +) +from core.rag.splitter.text_splitter import TextSplitter +from models.dataset import Dataset, DatasetProcessRule + + +class BaseIndexProcessor(ABC): + """Interface for extract files.""" + + @abstractmethod + def extract(self, extract_setting: ExtractSetting, **kwargs) -> list[Document]: + raise NotImplementedError + + @abstractmethod + def transform(self, documents: list[Document], **kwargs) -> list[Document]: + raise NotImplementedError + + @abstractmethod + def load(self, dataset: Dataset, documents: list[Document], with_keywords: bool = True, **kwargs): + raise NotImplementedError + + def clean(self, dataset: Dataset, node_ids: Optional[list[str]], with_keywords: bool = True, **kwargs): + raise NotImplementedError + + @abstractmethod + def retrieve( + self, + retrieval_method: str, + query: str, + dataset: Dataset, + top_k: int, + score_threshold: float, + reranking_model: dict, + ) -> list[Document]: + raise NotImplementedError + + def _get_splitter( + self, + processing_rule_mode: str, + max_tokens: int, + chunk_overlap: int, + separator: str, + embedding_model_instance: Optional[ModelInstance], + ) -> TextSplitter: + """ + Get the NodeParser object according to the processing rule. + """ + if processing_rule_mode in ["custom", "hierarchical"]: + # The user-defined segmentation rule + max_segmentation_tokens_length = dify_config.INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH + if max_tokens < 50 or max_tokens > max_segmentation_tokens_length: + raise ValueError(f"Custom segment length should be between 50 and {max_segmentation_tokens_length}.") + + if separator: + separator = separator.replace("\\n", "\n") + + character_splitter = FixedRecursiveCharacterTextSplitter.from_encoder( + chunk_size=max_tokens, + chunk_overlap=chunk_overlap, + fixed_separator=separator, + separators=["\n\n", "。", ". ", " ", ""], + embedding_model_instance=embedding_model_instance, + ) + else: + # Automatic segmentation + character_splitter = EnhanceRecursiveCharacterTextSplitter.from_encoder( + chunk_size=DatasetProcessRule.AUTOMATIC_RULES["segmentation"]["max_tokens"], + chunk_overlap=DatasetProcessRule.AUTOMATIC_RULES["segmentation"]["chunk_overlap"], + separators=["\n\n", "。", ". ", " ", ""], + embedding_model_instance=embedding_model_instance, + ) + + return character_splitter # type: ignore diff --git a/api/core/rag/index_processor/index_processor_factory.py b/api/core/rag/index_processor/index_processor_factory.py new file mode 100644 index 0000000000000000000000000000000000000000..c987edf342ab8ac6ad265f455a9a5d1d75c6838d --- /dev/null +++ b/api/core/rag/index_processor/index_processor_factory.py @@ -0,0 +1,29 @@ +"""Abstract interface for document loader implementations.""" + +from core.rag.index_processor.constant.index_type import IndexType +from core.rag.index_processor.index_processor_base import BaseIndexProcessor +from core.rag.index_processor.processor.paragraph_index_processor import ParagraphIndexProcessor +from core.rag.index_processor.processor.parent_child_index_processor import ParentChildIndexProcessor +from core.rag.index_processor.processor.qa_index_processor import QAIndexProcessor + + +class IndexProcessorFactory: + """IndexProcessorInit.""" + + def __init__(self, index_type: str | None): + self._index_type = index_type + + def init_index_processor(self) -> BaseIndexProcessor: + """Init index processor.""" + + if not self._index_type: + raise ValueError("Index type must be specified.") + + if self._index_type == IndexType.PARAGRAPH_INDEX: + return ParagraphIndexProcessor() + elif self._index_type == IndexType.QA_INDEX: + return QAIndexProcessor() + elif self._index_type == IndexType.PARENT_CHILD_INDEX: + return ParentChildIndexProcessor() + else: + raise ValueError(f"Index type {self._index_type} is not supported.") diff --git a/api/core/rag/index_processor/processor/__init__.py b/api/core/rag/index_processor/processor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/index_processor/processor/paragraph_index_processor.py b/api/core/rag/index_processor/processor/paragraph_index_processor.py new file mode 100644 index 0000000000000000000000000000000000000000..dca84b90416e0d0b1d0b3e36bc24482f2a236c5b --- /dev/null +++ b/api/core/rag/index_processor/processor/paragraph_index_processor.py @@ -0,0 +1,127 @@ +"""Paragraph index processor.""" + +import uuid +from typing import Optional + +from core.rag.cleaner.clean_processor import CleanProcessor +from core.rag.datasource.keyword.keyword_factory import Keyword +from core.rag.datasource.retrieval_service import RetrievalService +from core.rag.datasource.vdb.vector_factory import Vector +from core.rag.extractor.entity.extract_setting import ExtractSetting +from core.rag.extractor.extract_processor import ExtractProcessor +from core.rag.index_processor.index_processor_base import BaseIndexProcessor +from core.rag.models.document import Document +from core.tools.utils.text_processing_utils import remove_leading_symbols +from libs import helper +from models.dataset import Dataset, DatasetProcessRule +from services.entities.knowledge_entities.knowledge_entities import Rule + + +class ParagraphIndexProcessor(BaseIndexProcessor): + def extract(self, extract_setting: ExtractSetting, **kwargs) -> list[Document]: + text_docs = ExtractProcessor.extract( + extract_setting=extract_setting, + is_automatic=( + kwargs.get("process_rule_mode") == "automatic" or kwargs.get("process_rule_mode") == "hierarchical" + ), + ) + + return text_docs + + def transform(self, documents: list[Document], **kwargs) -> list[Document]: + process_rule = kwargs.get("process_rule") + if not process_rule: + raise ValueError("No process rule found.") + if process_rule.get("mode") == "automatic": + automatic_rule = DatasetProcessRule.AUTOMATIC_RULES + rules = Rule(**automatic_rule) + else: + if not process_rule.get("rules"): + raise ValueError("No rules found in process rule.") + rules = Rule(**process_rule.get("rules")) + # Split the text documents into nodes. + if not rules.segmentation: + raise ValueError("No segmentation found in rules.") + splitter = self._get_splitter( + processing_rule_mode=process_rule.get("mode"), + max_tokens=rules.segmentation.max_tokens, + chunk_overlap=rules.segmentation.chunk_overlap, + separator=rules.segmentation.separator, + embedding_model_instance=kwargs.get("embedding_model_instance"), + ) + all_documents = [] + for document in documents: + # document clean + document_text = CleanProcessor.clean(document.page_content, kwargs.get("process_rule", {})) + document.page_content = document_text + # parse document to nodes + document_nodes = splitter.split_documents([document]) + split_documents = [] + for document_node in document_nodes: + if document_node.page_content.strip(): + doc_id = str(uuid.uuid4()) + hash = helper.generate_text_hash(document_node.page_content) + if document_node.metadata is not None: + document_node.metadata["doc_id"] = doc_id + document_node.metadata["doc_hash"] = hash + # delete Splitter character + page_content = remove_leading_symbols(document_node.page_content).strip() + if len(page_content) > 0: + document_node.page_content = page_content + split_documents.append(document_node) + all_documents.extend(split_documents) + return all_documents + + def load(self, dataset: Dataset, documents: list[Document], with_keywords: bool = True, **kwargs): + if dataset.indexing_technique == "high_quality": + vector = Vector(dataset) + vector.create(documents) + if with_keywords: + keywords_list = kwargs.get("keywords_list") + keyword = Keyword(dataset) + if keywords_list and len(keywords_list) > 0: + keyword.add_texts(documents, keywords_list=keywords_list) + else: + keyword.add_texts(documents) + + def clean(self, dataset: Dataset, node_ids: Optional[list[str]], with_keywords: bool = True, **kwargs): + if dataset.indexing_technique == "high_quality": + vector = Vector(dataset) + if node_ids: + vector.delete_by_ids(node_ids) + else: + vector.delete() + if with_keywords: + keyword = Keyword(dataset) + if node_ids: + keyword.delete_by_ids(node_ids) + else: + keyword.delete() + + def retrieve( + self, + retrieval_method: str, + query: str, + dataset: Dataset, + top_k: int, + score_threshold: float, + reranking_model: dict, + ) -> list[Document]: + # Set search parameters. + results = RetrievalService.retrieve( + retrieval_method=retrieval_method, + dataset_id=dataset.id, + query=query, + top_k=top_k, + score_threshold=score_threshold, + reranking_model=reranking_model, + ) + # Organize results. + docs = [] + for result in results: + metadata = result.metadata + metadata["score"] = result.score + if result.score > score_threshold: + doc = Document(page_content=result.page_content, metadata=metadata) + docs.append(doc) + return docs diff --git a/api/core/rag/index_processor/processor/parent_child_index_processor.py b/api/core/rag/index_processor/processor/parent_child_index_processor.py new file mode 100644 index 0000000000000000000000000000000000000000..31401220818d20c9f3f89b45763ed58981f094a5 --- /dev/null +++ b/api/core/rag/index_processor/processor/parent_child_index_processor.py @@ -0,0 +1,200 @@ +"""Paragraph index processor.""" + +import uuid +from typing import Optional + +from configs import dify_config +from core.model_manager import ModelInstance +from core.rag.cleaner.clean_processor import CleanProcessor +from core.rag.datasource.retrieval_service import RetrievalService +from core.rag.datasource.vdb.vector_factory import Vector +from core.rag.extractor.entity.extract_setting import ExtractSetting +from core.rag.extractor.extract_processor import ExtractProcessor +from core.rag.index_processor.index_processor_base import BaseIndexProcessor +from core.rag.models.document import ChildDocument, Document +from extensions.ext_database import db +from libs import helper +from models.dataset import ChildChunk, Dataset, DocumentSegment +from services.entities.knowledge_entities.knowledge_entities import ParentMode, Rule + + +class ParentChildIndexProcessor(BaseIndexProcessor): + def extract(self, extract_setting: ExtractSetting, **kwargs) -> list[Document]: + text_docs = ExtractProcessor.extract( + extract_setting=extract_setting, + is_automatic=( + kwargs.get("process_rule_mode") == "automatic" or kwargs.get("process_rule_mode") == "hierarchical" + ), + ) + + return text_docs + + def transform(self, documents: list[Document], **kwargs) -> list[Document]: + process_rule = kwargs.get("process_rule") + if not process_rule: + raise ValueError("No process rule found.") + if not process_rule.get("rules"): + raise ValueError("No rules found in process rule.") + rules = Rule(**process_rule.get("rules")) + all_documents = [] # type: ignore + if rules.parent_mode == ParentMode.PARAGRAPH: + # Split the text documents into nodes. + splitter = self._get_splitter( + processing_rule_mode=process_rule.get("mode"), + max_tokens=rules.segmentation.max_tokens, + chunk_overlap=rules.segmentation.chunk_overlap, + separator=rules.segmentation.separator, + embedding_model_instance=kwargs.get("embedding_model_instance"), + ) + for document in documents: + # document clean + document_text = CleanProcessor.clean(document.page_content, process_rule) + document.page_content = document_text + # parse document to nodes + document_nodes = splitter.split_documents([document]) + split_documents = [] + for document_node in document_nodes: + if document_node.page_content.strip(): + doc_id = str(uuid.uuid4()) + hash = helper.generate_text_hash(document_node.page_content) + document_node.metadata["doc_id"] = doc_id + document_node.metadata["doc_hash"] = hash + # delete Splitter character + page_content = document_node.page_content + if page_content.startswith(".") or page_content.startswith("。"): + page_content = page_content[1:].strip() + else: + page_content = page_content + if len(page_content) > 0: + document_node.page_content = page_content + # parse document to child nodes + child_nodes = self._split_child_nodes( + document_node, rules, process_rule.get("mode"), kwargs.get("embedding_model_instance") + ) + document_node.children = child_nodes + split_documents.append(document_node) + all_documents.extend(split_documents) + elif rules.parent_mode == ParentMode.FULL_DOC: + page_content = "\n".join([document.page_content for document in documents]) + document = Document(page_content=page_content, metadata=documents[0].metadata) + # parse document to child nodes + child_nodes = self._split_child_nodes( + document, rules, process_rule.get("mode"), kwargs.get("embedding_model_instance") + ) + if kwargs.get("preview"): + if len(child_nodes) > dify_config.CHILD_CHUNKS_PREVIEW_NUMBER: + child_nodes = child_nodes[: dify_config.CHILD_CHUNKS_PREVIEW_NUMBER] + + document.children = child_nodes + doc_id = str(uuid.uuid4()) + hash = helper.generate_text_hash(document.page_content) + document.metadata["doc_id"] = doc_id + document.metadata["doc_hash"] = hash + all_documents.append(document) + + return all_documents + + def load(self, dataset: Dataset, documents: list[Document], with_keywords: bool = True, **kwargs): + if dataset.indexing_technique == "high_quality": + vector = Vector(dataset) + for document in documents: + child_documents = document.children + if child_documents: + formatted_child_documents = [ + Document(**child_document.model_dump()) for child_document in child_documents + ] + vector.create(formatted_child_documents) + + def clean(self, dataset: Dataset, node_ids: Optional[list[str]], with_keywords: bool = True, **kwargs): + # node_ids is segment's node_ids + if dataset.indexing_technique == "high_quality": + delete_child_chunks = kwargs.get("delete_child_chunks") or False + vector = Vector(dataset) + if node_ids: + child_node_ids = ( + db.session.query(ChildChunk.index_node_id) + .join(DocumentSegment, ChildChunk.segment_id == DocumentSegment.id) + .filter( + DocumentSegment.dataset_id == dataset.id, + DocumentSegment.index_node_id.in_(node_ids), + ChildChunk.dataset_id == dataset.id, + ) + .all() + ) + child_node_ids = [child_node_id[0] for child_node_id in child_node_ids] + vector.delete_by_ids(child_node_ids) + if delete_child_chunks: + db.session.query(ChildChunk).filter( + ChildChunk.dataset_id == dataset.id, ChildChunk.index_node_id.in_(child_node_ids) + ).delete() + db.session.commit() + else: + vector.delete() + + if delete_child_chunks: + db.session.query(ChildChunk).filter(ChildChunk.dataset_id == dataset.id).delete() + db.session.commit() + + def retrieve( + self, + retrieval_method: str, + query: str, + dataset: Dataset, + top_k: int, + score_threshold: float, + reranking_model: dict, + ) -> list[Document]: + # Set search parameters. + results = RetrievalService.retrieve( + retrieval_method=retrieval_method, + dataset_id=dataset.id, + query=query, + top_k=top_k, + score_threshold=score_threshold, + reranking_model=reranking_model, + ) + # Organize results. + docs = [] + for result in results: + metadata = result.metadata + metadata["score"] = result.score + if result.score > score_threshold: + doc = Document(page_content=result.page_content, metadata=metadata) + docs.append(doc) + return docs + + def _split_child_nodes( + self, + document_node: Document, + rules: Rule, + process_rule_mode: str, + embedding_model_instance: Optional[ModelInstance], + ) -> list[ChildDocument]: + if not rules.subchunk_segmentation: + raise ValueError("No subchunk segmentation found in rules.") + child_splitter = self._get_splitter( + processing_rule_mode=process_rule_mode, + max_tokens=rules.subchunk_segmentation.max_tokens, + chunk_overlap=rules.subchunk_segmentation.chunk_overlap, + separator=rules.subchunk_segmentation.separator, + embedding_model_instance=embedding_model_instance, + ) + # parse document to child nodes + child_nodes = [] + child_documents = child_splitter.split_documents([document_node]) + for child_document_node in child_documents: + if child_document_node.page_content.strip(): + doc_id = str(uuid.uuid4()) + hash = helper.generate_text_hash(child_document_node.page_content) + child_document = ChildDocument( + page_content=child_document_node.page_content, metadata=document_node.metadata + ) + child_document.metadata["doc_id"] = doc_id + child_document.metadata["doc_hash"] = hash + child_page_content = child_document.page_content + if child_page_content.startswith(".") or child_page_content.startswith("。"): + child_page_content = child_page_content[1:].strip() + if len(child_page_content) > 0: + child_document.page_content = child_page_content + child_nodes.append(child_document) + return child_nodes diff --git a/api/core/rag/index_processor/processor/qa_index_processor.py b/api/core/rag/index_processor/processor/qa_index_processor.py new file mode 100644 index 0000000000000000000000000000000000000000..0055625e136c798cdb51a908e4a2fe66a0904a6c --- /dev/null +++ b/api/core/rag/index_processor/processor/qa_index_processor.py @@ -0,0 +1,193 @@ +"""Paragraph index processor.""" + +import logging +import re +import threading +import uuid +from typing import Optional + +import pandas as pd +from flask import Flask, current_app +from werkzeug.datastructures import FileStorage + +from core.llm_generator.llm_generator import LLMGenerator +from core.rag.cleaner.clean_processor import CleanProcessor +from core.rag.datasource.retrieval_service import RetrievalService +from core.rag.datasource.vdb.vector_factory import Vector +from core.rag.extractor.entity.extract_setting import ExtractSetting +from core.rag.extractor.extract_processor import ExtractProcessor +from core.rag.index_processor.index_processor_base import BaseIndexProcessor +from core.rag.models.document import Document +from core.tools.utils.text_processing_utils import remove_leading_symbols +from libs import helper +from models.dataset import Dataset +from services.entities.knowledge_entities.knowledge_entities import Rule + + +class QAIndexProcessor(BaseIndexProcessor): + def extract(self, extract_setting: ExtractSetting, **kwargs) -> list[Document]: + text_docs = ExtractProcessor.extract( + extract_setting=extract_setting, + is_automatic=( + kwargs.get("process_rule_mode") == "automatic" or kwargs.get("process_rule_mode") == "hierarchical" + ), + ) + return text_docs + + def transform(self, documents: list[Document], **kwargs) -> list[Document]: + preview = kwargs.get("preview") + process_rule = kwargs.get("process_rule") + if not process_rule: + raise ValueError("No process rule found.") + if not process_rule.get("rules"): + raise ValueError("No rules found in process rule.") + rules = Rule(**process_rule.get("rules")) + splitter = self._get_splitter( + processing_rule_mode=process_rule.get("mode"), + max_tokens=rules.segmentation.max_tokens if rules.segmentation else 0, + chunk_overlap=rules.segmentation.chunk_overlap if rules.segmentation else 0, + separator=rules.segmentation.separator if rules.segmentation else "", + embedding_model_instance=kwargs.get("embedding_model_instance"), + ) + + # Split the text documents into nodes. + all_documents: list[Document] = [] + all_qa_documents: list[Document] = [] + for document in documents: + # document clean + document_text = CleanProcessor.clean(document.page_content, kwargs.get("process_rule") or {}) + document.page_content = document_text + + # parse document to nodes + document_nodes = splitter.split_documents([document]) + split_documents = [] + for document_node in document_nodes: + if document_node.page_content.strip(): + doc_id = str(uuid.uuid4()) + hash = helper.generate_text_hash(document_node.page_content) + if document_node.metadata is not None: + document_node.metadata["doc_id"] = doc_id + document_node.metadata["doc_hash"] = hash + # delete Splitter character + page_content = document_node.page_content + document_node.page_content = remove_leading_symbols(page_content) + split_documents.append(document_node) + all_documents.extend(split_documents) + if preview: + self._format_qa_document( + current_app._get_current_object(), # type: ignore + kwargs.get("tenant_id"), # type: ignore + all_documents[0], + all_qa_documents, + kwargs.get("doc_language", "English"), + ) + else: + for i in range(0, len(all_documents), 10): + threads = [] + sub_documents = all_documents[i : i + 10] + for doc in sub_documents: + document_format_thread = threading.Thread( + target=self._format_qa_document, + kwargs={ + "flask_app": current_app._get_current_object(), # type: ignore + "tenant_id": kwargs.get("tenant_id"), # type: ignore + "document_node": doc, + "all_qa_documents": all_qa_documents, + "document_language": kwargs.get("doc_language", "English"), + }, + ) + threads.append(document_format_thread) + document_format_thread.start() + for thread in threads: + thread.join() + return all_qa_documents + + def format_by_template(self, file: FileStorage, **kwargs) -> list[Document]: + # check file type + if not file.filename.endswith(".csv"): + raise ValueError("Invalid file type. Only CSV files are allowed") + + try: + # Skip the first row + df = pd.read_csv(file) + text_docs = [] + for index, row in df.iterrows(): + data = Document(page_content=row.iloc[0], metadata={"answer": row.iloc[1]}) + text_docs.append(data) + if len(text_docs) == 0: + raise ValueError("The CSV file is empty.") + + except Exception as e: + raise ValueError(str(e)) + return text_docs + + def load(self, dataset: Dataset, documents: list[Document], with_keywords: bool = True, **kwargs): + if dataset.indexing_technique == "high_quality": + vector = Vector(dataset) + vector.create(documents) + + def clean(self, dataset: Dataset, node_ids: Optional[list[str]], with_keywords: bool = True, **kwargs): + vector = Vector(dataset) + if node_ids: + vector.delete_by_ids(node_ids) + else: + vector.delete() + + def retrieve( + self, + retrieval_method: str, + query: str, + dataset: Dataset, + top_k: int, + score_threshold: float, + reranking_model: dict, + ): + # Set search parameters. + results = RetrievalService.retrieve( + retrieval_method=retrieval_method, + dataset_id=dataset.id, + query=query, + top_k=top_k, + score_threshold=score_threshold, + reranking_model=reranking_model, + ) + # Organize results. + docs = [] + for result in results: + metadata = result.metadata + metadata["score"] = result.score + if result.score > score_threshold: + doc = Document(page_content=result.page_content, metadata=metadata) + docs.append(doc) + return docs + + def _format_qa_document(self, flask_app: Flask, tenant_id: str, document_node, all_qa_documents, document_language): + format_documents = [] + if document_node.page_content is None or not document_node.page_content.strip(): + return + with flask_app.app_context(): + try: + # qa model document + response = LLMGenerator.generate_qa_document(tenant_id, document_node.page_content, document_language) + document_qa_list = self._format_split_text(response) + qa_documents = [] + for result in document_qa_list: + qa_document = Document(page_content=result["question"], metadata=document_node.metadata.copy()) + if qa_document.metadata is not None: + doc_id = str(uuid.uuid4()) + hash = helper.generate_text_hash(result["question"]) + qa_document.metadata["answer"] = result["answer"] + qa_document.metadata["doc_id"] = doc_id + qa_document.metadata["doc_hash"] = hash + qa_documents.append(qa_document) + format_documents.extend(qa_documents) + except Exception as e: + logging.exception("Failed to format qa document") + + all_qa_documents.extend(format_documents) + + def _format_split_text(self, text): + regex = r"Q\d+:\s*(.*?)\s*A\d+:\s*([\s\S]*?)(?=Q\d+:|$)" + matches = re.findall(regex, text, re.UNICODE) + + return [{"question": q, "answer": re.sub(r"\n\s*", "\n", a.strip())} for q, a in matches if q and a] diff --git a/api/core/rag/models/__init__.py b/api/core/rag/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/models/document.py b/api/core/rag/models/document.py new file mode 100644 index 0000000000000000000000000000000000000000..421cdc05df7cc06b5966a93276faf77d8e349cd7 --- /dev/null +++ b/api/core/rag/models/document.py @@ -0,0 +1,94 @@ +from abc import ABC, abstractmethod +from collections.abc import Sequence +from typing import Any, Optional + +from pydantic import BaseModel + + +class ChildDocument(BaseModel): + """Class for storing a piece of text and associated metadata.""" + + page_content: str + + vector: Optional[list[float]] = None + + """Arbitrary metadata about the page content (e.g., source, relationships to other + documents, etc.). + """ + metadata: dict = {} + + +class Document(BaseModel): + """Class for storing a piece of text and associated metadata.""" + + page_content: str + + vector: Optional[list[float]] = None + + """Arbitrary metadata about the page content (e.g., source, relationships to other + documents, etc.). + """ + metadata: dict = {} + + provider: Optional[str] = "dify" + + children: Optional[list[ChildDocument]] = None + + +class BaseDocumentTransformer(ABC): + """Abstract base class for document transformation systems. + + A document transformation system takes a sequence of Documents and returns a + sequence of transformed Documents. + + Example: + .. code-block:: python + + class EmbeddingsRedundantFilter(BaseDocumentTransformer, BaseModel): + embeddings: Embeddings + similarity_fn: Callable = cosine_similarity + similarity_threshold: float = 0.95 + + class Config: + arbitrary_types_allowed = True + + def transform_documents( + self, documents: Sequence[Document], **kwargs: Any + ) -> Sequence[Document]: + stateful_documents = get_stateful_documents(documents) + embedded_documents = _get_embeddings_from_stateful_docs( + self.embeddings, stateful_documents + ) + included_idxs = _filter_similar_embeddings( + embedded_documents, self.similarity_fn, self.similarity_threshold + ) + return [stateful_documents[i] for i in sorted(included_idxs)] + + async def atransform_documents( + self, documents: Sequence[Document], **kwargs: Any + ) -> Sequence[Document]: + raise NotImplementedError + + """ + + @abstractmethod + def transform_documents(self, documents: Sequence[Document], **kwargs: Any) -> Sequence[Document]: + """Transform a list of documents. + + Args: + documents: A sequence of Documents to be transformed. + + Returns: + A list of transformed Documents. + """ + + @abstractmethod + async def atransform_documents(self, documents: Sequence[Document], **kwargs: Any) -> Sequence[Document]: + """Asynchronously transform a list of documents. + + Args: + documents: A sequence of Documents to be transformed. + + Returns: + A list of transformed Documents. + """ diff --git a/api/core/rag/rerank/__init__.py b/api/core/rag/rerank/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/rerank/entity/weight.py b/api/core/rag/rerank/entity/weight.py new file mode 100644 index 0000000000000000000000000000000000000000..6dbbad2f8da61a75ef35ab2c3e89fc49b1920fe7 --- /dev/null +++ b/api/core/rag/rerank/entity/weight.py @@ -0,0 +1,21 @@ +from pydantic import BaseModel + + +class VectorSetting(BaseModel): + vector_weight: float + + embedding_provider_name: str + + embedding_model_name: str + + +class KeywordSetting(BaseModel): + keyword_weight: float + + +class Weights(BaseModel): + """Model for weighted rerank.""" + + vector_setting: VectorSetting + + keyword_setting: KeywordSetting diff --git a/api/core/rag/rerank/rerank_base.py b/api/core/rag/rerank/rerank_base.py new file mode 100644 index 0000000000000000000000000000000000000000..818b04b2ffc196f4ff69bb4e28a69fd8dbbdf220 --- /dev/null +++ b/api/core/rag/rerank/rerank_base.py @@ -0,0 +1,26 @@ +from abc import ABC, abstractmethod +from typing import Optional + +from core.rag.models.document import Document + + +class BaseRerankRunner(ABC): + @abstractmethod + def run( + self, + query: str, + documents: list[Document], + score_threshold: Optional[float] = None, + top_n: Optional[int] = None, + user: Optional[str] = None, + ) -> list[Document]: + """ + Run rerank model + :param query: search query + :param documents: documents for reranking + :param score_threshold: score threshold + :param top_n: top n + :param user: unique user id if needed + :return: + """ + raise NotImplementedError diff --git a/api/core/rag/rerank/rerank_factory.py b/api/core/rag/rerank/rerank_factory.py new file mode 100644 index 0000000000000000000000000000000000000000..1a3cf8573631f26149f133a1de2ebe1d90b3519d --- /dev/null +++ b/api/core/rag/rerank/rerank_factory.py @@ -0,0 +1,16 @@ +from core.rag.rerank.rerank_base import BaseRerankRunner +from core.rag.rerank.rerank_model import RerankModelRunner +from core.rag.rerank.rerank_type import RerankMode +from core.rag.rerank.weight_rerank import WeightRerankRunner + + +class RerankRunnerFactory: + @staticmethod + def create_rerank_runner(runner_type: str, *args, **kwargs) -> BaseRerankRunner: + match runner_type: + case RerankMode.RERANKING_MODEL.value: + return RerankModelRunner(*args, **kwargs) + case RerankMode.WEIGHTED_SCORE.value: + return WeightRerankRunner(*args, **kwargs) + case _: + raise ValueError(f"Unknown runner type: {runner_type}") diff --git a/api/core/rag/rerank/rerank_model.py b/api/core/rag/rerank/rerank_model.py new file mode 100644 index 0000000000000000000000000000000000000000..ac7a3f8bb857e43e5873cbbac683578de359df88 --- /dev/null +++ b/api/core/rag/rerank/rerank_model.py @@ -0,0 +1,65 @@ +from typing import Optional + +from core.model_manager import ModelInstance +from core.rag.models.document import Document +from core.rag.rerank.rerank_base import BaseRerankRunner + + +class RerankModelRunner(BaseRerankRunner): + def __init__(self, rerank_model_instance: ModelInstance) -> None: + self.rerank_model_instance = rerank_model_instance + + def run( + self, + query: str, + documents: list[Document], + score_threshold: Optional[float] = None, + top_n: Optional[int] = None, + user: Optional[str] = None, + ) -> list[Document]: + """ + Run rerank model + :param query: search query + :param documents: documents for reranking + :param score_threshold: score threshold + :param top_n: top n + :param user: unique user id if needed + :return: + """ + docs = [] + doc_ids = set() + unique_documents = [] + for document in documents: + if ( + document.provider == "dify" + and document.metadata is not None + and document.metadata["doc_id"] not in doc_ids + ): + doc_ids.add(document.metadata["doc_id"]) + docs.append(document.page_content) + unique_documents.append(document) + elif document.provider == "external": + if document not in unique_documents: + docs.append(document.page_content) + unique_documents.append(document) + + documents = unique_documents + + rerank_result = self.rerank_model_instance.invoke_rerank( + query=query, docs=docs, score_threshold=score_threshold, top_n=top_n, user=user + ) + + rerank_documents = [] + + for result in rerank_result.docs: + # format document + rerank_document = Document( + page_content=result.text, + metadata=documents[result.index].metadata, + provider=documents[result.index].provider, + ) + if rerank_document.metadata is not None: + rerank_document.metadata["score"] = result.score + rerank_documents.append(rerank_document) + + return rerank_documents diff --git a/api/core/rag/rerank/rerank_type.py b/api/core/rag/rerank/rerank_type.py new file mode 100644 index 0000000000000000000000000000000000000000..b2d1314654044286ad5484c70d0c2c47d143e60c --- /dev/null +++ b/api/core/rag/rerank/rerank_type.py @@ -0,0 +1,6 @@ +from enum import StrEnum + + +class RerankMode(StrEnum): + RERANKING_MODEL = "reranking_model" + WEIGHTED_SCORE = "weighted_score" diff --git a/api/core/rag/rerank/weight_rerank.py b/api/core/rag/rerank/weight_rerank.py new file mode 100644 index 0000000000000000000000000000000000000000..cbc96037bf2cc0a47b7c4f105fad8a259553147a --- /dev/null +++ b/api/core/rag/rerank/weight_rerank.py @@ -0,0 +1,185 @@ +import math +from collections import Counter +from typing import Optional + +import numpy as np + +from core.model_manager import ModelManager +from core.model_runtime.entities.model_entities import ModelType +from core.rag.datasource.keyword.jieba.jieba_keyword_table_handler import JiebaKeywordTableHandler +from core.rag.embedding.cached_embedding import CacheEmbedding +from core.rag.models.document import Document +from core.rag.rerank.entity.weight import VectorSetting, Weights +from core.rag.rerank.rerank_base import BaseRerankRunner + + +class WeightRerankRunner(BaseRerankRunner): + def __init__(self, tenant_id: str, weights: Weights) -> None: + self.tenant_id = tenant_id + self.weights = weights + + def run( + self, + query: str, + documents: list[Document], + score_threshold: Optional[float] = None, + top_n: Optional[int] = None, + user: Optional[str] = None, + ) -> list[Document]: + """ + Run rerank model + :param query: search query + :param documents: documents for reranking + :param score_threshold: score threshold + :param top_n: top n + :param user: unique user id if needed + + :return: + """ + unique_documents = [] + doc_ids = set() + for document in documents: + if document.metadata is not None and document.metadata["doc_id"] not in doc_ids: + doc_ids.add(document.metadata["doc_id"]) + unique_documents.append(document) + + documents = unique_documents + + query_scores = self._calculate_keyword_score(query, documents) + query_vector_scores = self._calculate_cosine(self.tenant_id, query, documents, self.weights.vector_setting) + + rerank_documents = [] + for document, query_score, query_vector_score in zip(documents, query_scores, query_vector_scores): + score = ( + self.weights.vector_setting.vector_weight * query_vector_score + + self.weights.keyword_setting.keyword_weight * query_score + ) + if score_threshold and score < score_threshold: + continue + if document.metadata is not None: + document.metadata["score"] = score + rerank_documents.append(document) + + rerank_documents.sort(key=lambda x: x.metadata["score"] if x.metadata else 0, reverse=True) + return rerank_documents[:top_n] if top_n else rerank_documents + + def _calculate_keyword_score(self, query: str, documents: list[Document]) -> list[float]: + """ + Calculate BM25 scores + :param query: search query + :param documents: documents for reranking + + :return: + """ + keyword_table_handler = JiebaKeywordTableHandler() + query_keywords = keyword_table_handler.extract_keywords(query, None) + documents_keywords = [] + for document in documents: + # get the document keywords + document_keywords = keyword_table_handler.extract_keywords(document.page_content, None) + if document.metadata is not None: + document.metadata["keywords"] = document_keywords + documents_keywords.append(document_keywords) + + # Counter query keywords(TF) + query_keyword_counts = Counter(query_keywords) + + # total documents + total_documents = len(documents) + + # calculate all documents' keywords IDF + all_keywords = set() + for document_keywords in documents_keywords: + all_keywords.update(document_keywords) + + keyword_idf = {} + for keyword in all_keywords: + # calculate include query keywords' documents + doc_count_containing_keyword = sum(1 for doc_keywords in documents_keywords if keyword in doc_keywords) + # IDF + keyword_idf[keyword] = math.log((1 + total_documents) / (1 + doc_count_containing_keyword)) + 1 + + query_tfidf = {} + + for keyword, count in query_keyword_counts.items(): + tf = count + idf = keyword_idf.get(keyword, 0) + query_tfidf[keyword] = tf * idf + + # calculate all documents' TF-IDF + documents_tfidf = [] + for document_keywords in documents_keywords: + document_keyword_counts = Counter(document_keywords) + document_tfidf = {} + for keyword, count in document_keyword_counts.items(): + tf = count + idf = keyword_idf.get(keyword, 0) + document_tfidf[keyword] = tf * idf + documents_tfidf.append(document_tfidf) + + def cosine_similarity(vec1, vec2): + intersection = set(vec1.keys()) & set(vec2.keys()) + numerator = sum(vec1[x] * vec2[x] for x in intersection) + + sum1 = sum(vec1[x] ** 2 for x in vec1) + sum2 = sum(vec2[x] ** 2 for x in vec2) + denominator = math.sqrt(sum1) * math.sqrt(sum2) + + if not denominator: + return 0.0 + else: + return float(numerator) / denominator + + similarities = [] + for document_tfidf in documents_tfidf: + similarity = cosine_similarity(query_tfidf, document_tfidf) + similarities.append(similarity) + + # for idx, similarity in enumerate(similarities): + # print(f"Document {idx + 1} similarity: {similarity}") + + return similarities + + def _calculate_cosine( + self, tenant_id: str, query: str, documents: list[Document], vector_setting: VectorSetting + ) -> list[float]: + """ + Calculate Cosine scores + :param query: search query + :param documents: documents for reranking + + :return: + """ + query_vector_scores = [] + + model_manager = ModelManager() + + embedding_model = model_manager.get_model_instance( + tenant_id=tenant_id, + provider=vector_setting.embedding_provider_name, + model_type=ModelType.TEXT_EMBEDDING, + model=vector_setting.embedding_model_name, + ) + cache_embedding = CacheEmbedding(embedding_model) + query_vector = cache_embedding.embed_query(query) + for document in documents: + # calculate cosine similarity + if document.metadata and "score" in document.metadata: + query_vector_scores.append(document.metadata["score"]) + else: + # transform to NumPy + vec1 = np.array(query_vector) + vec2 = np.array(document.vector) + + # calculate dot product + dot_product = np.dot(vec1, vec2) + + # calculate norm + norm_vec1 = np.linalg.norm(vec1) + norm_vec2 = np.linalg.norm(vec2) + + # calculate cosine similarity + cosine_sim = dot_product / (norm_vec1 * norm_vec2) + query_vector_scores.append(cosine_sim) + + return query_vector_scores diff --git a/api/core/rag/retrieval/__init__.py b/api/core/rag/retrieval/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/retrieval/dataset_retrieval.py b/api/core/rag/retrieval/dataset_retrieval.py new file mode 100644 index 0000000000000000000000000000000000000000..e1d36aad1fa5d7ae03388a092ff5651ae0524280 --- /dev/null +++ b/api/core/rag/retrieval/dataset_retrieval.py @@ -0,0 +1,711 @@ +import math +import threading +from collections import Counter +from typing import Any, Optional, cast + +from flask import Flask, current_app + +from core.app.app_config.entities import DatasetEntity, DatasetRetrieveConfigEntity +from core.app.entities.app_invoke_entities import InvokeFrom, ModelConfigWithCredentialsEntity +from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler +from core.entities.agent_entities import PlanningStrategy +from core.memory.token_buffer_memory import TokenBufferMemory +from core.model_manager import ModelInstance, ModelManager +from core.model_runtime.entities.message_entities import PromptMessageTool +from core.model_runtime.entities.model_entities import ModelFeature, ModelType +from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel +from core.ops.entities.trace_entity import TraceTaskName +from core.ops.ops_trace_manager import TraceQueueManager, TraceTask +from core.ops.utils import measure_time +from core.rag.data_post_processor.data_post_processor import DataPostProcessor +from core.rag.datasource.keyword.jieba.jieba_keyword_table_handler import JiebaKeywordTableHandler +from core.rag.datasource.retrieval_service import RetrievalService +from core.rag.entities.context_entities import DocumentContext +from core.rag.models.document import Document +from core.rag.rerank.rerank_type import RerankMode +from core.rag.retrieval.retrieval_methods import RetrievalMethod +from core.rag.retrieval.router.multi_dataset_function_call_router import FunctionCallMultiDatasetRouter +from core.rag.retrieval.router.multi_dataset_react_route import ReactMultiDatasetRouter +from core.tools.tool.dataset_retriever.dataset_multi_retriever_tool import DatasetMultiRetrieverTool +from core.tools.tool.dataset_retriever.dataset_retriever_base_tool import DatasetRetrieverBaseTool +from core.tools.tool.dataset_retriever.dataset_retriever_tool import DatasetRetrieverTool +from extensions.ext_database import db +from models.dataset import Dataset, DatasetQuery, DocumentSegment +from models.dataset import Document as DatasetDocument +from services.external_knowledge_service import ExternalDatasetService + +default_retrieval_model: dict[str, Any] = { + "search_method": RetrievalMethod.SEMANTIC_SEARCH.value, + "reranking_enable": False, + "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""}, + "top_k": 2, + "score_threshold_enabled": False, +} + + +class DatasetRetrieval: + def __init__(self, application_generate_entity=None): + self.application_generate_entity = application_generate_entity + + def retrieve( + self, + app_id: str, + user_id: str, + tenant_id: str, + model_config: ModelConfigWithCredentialsEntity, + config: DatasetEntity, + query: str, + invoke_from: InvokeFrom, + show_retrieve_source: bool, + hit_callback: DatasetIndexToolCallbackHandler, + message_id: str, + memory: Optional[TokenBufferMemory] = None, + ) -> Optional[str]: + """ + Retrieve dataset. + :param app_id: app_id + :param user_id: user_id + :param tenant_id: tenant id + :param model_config: model config + :param config: dataset config + :param query: query + :param invoke_from: invoke from + :param show_retrieve_source: show retrieve source + :param hit_callback: hit callback + :param message_id: message id + :param memory: memory + :return: + """ + dataset_ids = config.dataset_ids + if len(dataset_ids) == 0: + return None + retrieve_config = config.retrieve_config + + # check model is support tool calling + model_type_instance = model_config.provider_model_bundle.model_type_instance + model_type_instance = cast(LargeLanguageModel, model_type_instance) + + model_manager = ModelManager() + model_instance = model_manager.get_model_instance( + tenant_id=tenant_id, model_type=ModelType.LLM, provider=model_config.provider, model=model_config.model + ) + + # get model schema + model_schema = model_type_instance.get_model_schema( + model=model_config.model, credentials=model_config.credentials + ) + + if not model_schema: + return None + + planning_strategy = PlanningStrategy.REACT_ROUTER + features = model_schema.features + if features: + if ModelFeature.TOOL_CALL in features or ModelFeature.MULTI_TOOL_CALL in features: + planning_strategy = PlanningStrategy.ROUTER + available_datasets = [] + for dataset_id in dataset_ids: + # get dataset from dataset id + dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() + + # pass if dataset is not available + if not dataset: + continue + + # pass if dataset is not available + if dataset and dataset.available_document_count == 0 and dataset.provider != "external": + continue + + available_datasets.append(dataset) + all_documents = [] + user_from = "account" if invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER} else "end_user" + if retrieve_config.retrieve_strategy == DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE: + all_documents = self.single_retrieve( + app_id, + tenant_id, + user_id, + user_from, + available_datasets, + query, + model_instance, + model_config, + planning_strategy, + message_id, + ) + elif retrieve_config.retrieve_strategy == DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE: + all_documents = self.multiple_retrieve( + app_id, + tenant_id, + user_id, + user_from, + available_datasets, + query, + retrieve_config.top_k or 0, + retrieve_config.score_threshold or 0, + retrieve_config.rerank_mode or "reranking_model", + retrieve_config.reranking_model, + retrieve_config.weights, + retrieve_config.reranking_enabled or True, + message_id, + ) + + dify_documents = [item for item in all_documents if item.provider == "dify"] + external_documents = [item for item in all_documents if item.provider == "external"] + document_context_list = [] + retrieval_resource_list = [] + # deal with external documents + for item in external_documents: + document_context_list.append(DocumentContext(content=item.page_content, score=item.metadata.get("score"))) + source = { + "dataset_id": item.metadata.get("dataset_id"), + "dataset_name": item.metadata.get("dataset_name"), + "document_name": item.metadata.get("title"), + "data_source_type": "external", + "retriever_from": invoke_from.to_source(), + "score": item.metadata.get("score"), + "content": item.page_content, + } + retrieval_resource_list.append(source) + # deal with dify documents + if dify_documents: + records = RetrievalService.format_retrieval_documents(dify_documents) + if records: + for record in records: + segment = record.segment + if segment.answer: + document_context_list.append( + DocumentContext( + content=f"question:{segment.get_sign_content()} answer:{segment.answer}", + score=record.score, + ) + ) + else: + document_context_list.append( + DocumentContext( + content=segment.get_sign_content(), + score=record.score, + ) + ) + if show_retrieve_source: + for record in records: + segment = record.segment + dataset = Dataset.query.filter_by(id=segment.dataset_id).first() + document = DatasetDocument.query.filter( + DatasetDocument.id == segment.document_id, + DatasetDocument.enabled == True, + DatasetDocument.archived == False, + ).first() + if dataset and document: + source = { + "dataset_id": dataset.id, + "dataset_name": dataset.name, + "document_id": document.id, + "document_name": document.name, + "data_source_type": document.data_source_type, + "segment_id": segment.id, + "retriever_from": invoke_from.to_source(), + "score": record.score or 0.0, + } + + if invoke_from.to_source() == "dev": + source["hit_count"] = segment.hit_count + source["word_count"] = segment.word_count + source["segment_position"] = segment.position + source["index_node_hash"] = segment.index_node_hash + if segment.answer: + source["content"] = f"question:{segment.content} \nanswer:{segment.answer}" + else: + source["content"] = segment.content + retrieval_resource_list.append(source) + if hit_callback and retrieval_resource_list: + retrieval_resource_list = sorted(retrieval_resource_list, key=lambda x: x.get("score") or 0.0, reverse=True) + for position, item in enumerate(retrieval_resource_list, start=1): + item["position"] = position + hit_callback.return_retriever_resource_info(retrieval_resource_list) + if document_context_list: + document_context_list = sorted(document_context_list, key=lambda x: x.score or 0.0, reverse=True) + return str("\n".join([document_context.content for document_context in document_context_list])) + return "" + + def single_retrieve( + self, + app_id: str, + tenant_id: str, + user_id: str, + user_from: str, + available_datasets: list, + query: str, + model_instance: ModelInstance, + model_config: ModelConfigWithCredentialsEntity, + planning_strategy: PlanningStrategy, + message_id: Optional[str] = None, + ): + tools = [] + for dataset in available_datasets: + description = dataset.description + if not description: + description = "useful for when you want to answer queries about the " + dataset.name + + description = description.replace("\n", "").replace("\r", "") + message_tool = PromptMessageTool( + name=dataset.id, + description=description, + parameters={ + "type": "object", + "properties": {}, + "required": [], + }, + ) + tools.append(message_tool) + dataset_id = None + if planning_strategy == PlanningStrategy.REACT_ROUTER: + react_multi_dataset_router = ReactMultiDatasetRouter() + dataset_id = react_multi_dataset_router.invoke( + query, tools, model_config, model_instance, user_id, tenant_id + ) + + elif planning_strategy == PlanningStrategy.ROUTER: + function_call_router = FunctionCallMultiDatasetRouter() + dataset_id = function_call_router.invoke(query, tools, model_config, model_instance) + + if dataset_id: + # get retrieval model config + dataset = db.session.query(Dataset).filter(Dataset.id == dataset_id).first() + if dataset: + results = [] + if dataset.provider == "external": + external_documents = ExternalDatasetService.fetch_external_knowledge_retrieval( + tenant_id=dataset.tenant_id, + dataset_id=dataset_id, + query=query, + external_retrieval_parameters=dataset.retrieval_model, + ) + for external_document in external_documents: + document = Document( + page_content=external_document.get("content"), + metadata=external_document.get("metadata"), + provider="external", + ) + if document.metadata is not None: + document.metadata["score"] = external_document.get("score") + document.metadata["title"] = external_document.get("title") + document.metadata["dataset_id"] = dataset_id + document.metadata["dataset_name"] = dataset.name + results.append(document) + else: + retrieval_model_config = dataset.retrieval_model or default_retrieval_model + + # get top k + top_k = retrieval_model_config["top_k"] + # get retrieval method + if dataset.indexing_technique == "economy": + retrieval_method = "keyword_search" + else: + retrieval_method = retrieval_model_config["search_method"] + # get reranking model + reranking_model = ( + retrieval_model_config["reranking_model"] + if retrieval_model_config["reranking_enable"] + else None + ) + # get score threshold + score_threshold = 0.0 + score_threshold_enabled = retrieval_model_config.get("score_threshold_enabled") + if score_threshold_enabled: + score_threshold = retrieval_model_config.get("score_threshold", 0.0) + + with measure_time() as timer: + results = RetrievalService.retrieve( + retrieval_method=retrieval_method, + dataset_id=dataset.id, + query=query, + top_k=top_k, + score_threshold=score_threshold, + reranking_model=reranking_model, + reranking_mode=retrieval_model_config.get("reranking_mode", "reranking_model"), + weights=retrieval_model_config.get("weights", None), + ) + self._on_query(query, [dataset_id], app_id, user_from, user_id) + + if results: + self._on_retrieval_end(results, message_id, timer) + + return results + return [] + + def multiple_retrieve( + self, + app_id: str, + tenant_id: str, + user_id: str, + user_from: str, + available_datasets: list, + query: str, + top_k: int, + score_threshold: float, + reranking_mode: str, + reranking_model: Optional[dict] = None, + weights: Optional[dict[str, Any]] = None, + reranking_enable: bool = True, + message_id: Optional[str] = None, + ): + if not available_datasets: + return [] + threads = [] + all_documents: list[Document] = [] + dataset_ids = [dataset.id for dataset in available_datasets] + index_type_check = all( + item.indexing_technique == available_datasets[0].indexing_technique for item in available_datasets + ) + if not index_type_check and (not reranking_enable or reranking_mode != RerankMode.RERANKING_MODEL): + raise ValueError( + "The configured knowledge base list have different indexing technique, please set reranking model." + ) + index_type = available_datasets[0].indexing_technique + if index_type == "high_quality": + embedding_model_check = all( + item.embedding_model == available_datasets[0].embedding_model for item in available_datasets + ) + embedding_model_provider_check = all( + item.embedding_model_provider == available_datasets[0].embedding_model_provider + for item in available_datasets + ) + if ( + reranking_enable + and reranking_mode == "weighted_score" + and (not embedding_model_check or not embedding_model_provider_check) + ): + raise ValueError( + "The configured knowledge base list have different embedding model, please set reranking model." + ) + if reranking_enable and reranking_mode == RerankMode.WEIGHTED_SCORE: + if weights is not None: + weights["vector_setting"]["embedding_provider_name"] = available_datasets[ + 0 + ].embedding_model_provider + weights["vector_setting"]["embedding_model_name"] = available_datasets[0].embedding_model + + for dataset in available_datasets: + index_type = dataset.indexing_technique + retrieval_thread = threading.Thread( + target=self._retriever, + kwargs={ + "flask_app": current_app._get_current_object(), # type: ignore + "dataset_id": dataset.id, + "query": query, + "top_k": top_k, + "all_documents": all_documents, + }, + ) + threads.append(retrieval_thread) + retrieval_thread.start() + for thread in threads: + thread.join() + + with measure_time() as timer: + if reranking_enable: + # do rerank for searched documents + data_post_processor = DataPostProcessor(tenant_id, reranking_mode, reranking_model, weights, False) + + all_documents = data_post_processor.invoke( + query=query, documents=all_documents, score_threshold=score_threshold, top_n=top_k + ) + else: + if index_type == "economy": + all_documents = self.calculate_keyword_score(query, all_documents, top_k) + elif index_type == "high_quality": + all_documents = self.calculate_vector_score(all_documents, top_k, score_threshold) + + self._on_query(query, dataset_ids, app_id, user_from, user_id) + + if all_documents: + self._on_retrieval_end(all_documents, message_id, timer) + + return all_documents + + def _on_retrieval_end( + self, documents: list[Document], message_id: Optional[str] = None, timer: Optional[dict] = None + ) -> None: + """Handle retrieval end.""" + dify_documents = [document for document in documents if document.provider == "dify"] + for document in dify_documents: + if document.metadata is not None: + query = db.session.query(DocumentSegment).filter( + DocumentSegment.index_node_id == document.metadata["doc_id"] + ) + + # if 'dataset_id' in document.metadata: + if "dataset_id" in document.metadata: + query = query.filter(DocumentSegment.dataset_id == document.metadata["dataset_id"]) + + # add hit count to document segment + query.update({DocumentSegment.hit_count: DocumentSegment.hit_count + 1}, synchronize_session=False) + + db.session.commit() + + # get tracing instance + trace_manager: Optional[TraceQueueManager] = ( + self.application_generate_entity.trace_manager if self.application_generate_entity else None + ) + if trace_manager: + trace_manager.add_trace_task( + TraceTask( + TraceTaskName.DATASET_RETRIEVAL_TRACE, message_id=message_id, documents=documents, timer=timer + ) + ) + + def _on_query(self, query: str, dataset_ids: list[str], app_id: str, user_from: str, user_id: str) -> None: + """ + Handle query. + """ + if not query: + return + dataset_queries = [] + for dataset_id in dataset_ids: + dataset_query = DatasetQuery( + dataset_id=dataset_id, + content=query, + source="app", + source_app_id=app_id, + created_by_role=user_from, + created_by=user_id, + ) + dataset_queries.append(dataset_query) + if dataset_queries: + db.session.add_all(dataset_queries) + db.session.commit() + + def _retriever(self, flask_app: Flask, dataset_id: str, query: str, top_k: int, all_documents: list): + with flask_app.app_context(): + dataset = db.session.query(Dataset).filter(Dataset.id == dataset_id).first() + + if not dataset: + return [] + + if dataset.provider == "external": + external_documents = ExternalDatasetService.fetch_external_knowledge_retrieval( + tenant_id=dataset.tenant_id, + dataset_id=dataset_id, + query=query, + external_retrieval_parameters=dataset.retrieval_model, + ) + for external_document in external_documents: + document = Document( + page_content=external_document.get("content"), + metadata=external_document.get("metadata"), + provider="external", + ) + if document.metadata is not None: + document.metadata["score"] = external_document.get("score") + document.metadata["title"] = external_document.get("title") + document.metadata["dataset_id"] = dataset_id + document.metadata["dataset_name"] = dataset.name + all_documents.append(document) + else: + # get retrieval model , if the model is not setting , using default + retrieval_model = dataset.retrieval_model or default_retrieval_model + + if dataset.indexing_technique == "economy": + # use keyword table query + documents = RetrievalService.retrieve( + retrieval_method="keyword_search", dataset_id=dataset.id, query=query, top_k=top_k + ) + if documents: + all_documents.extend(documents) + else: + if top_k > 0: + # retrieval source + documents = RetrievalService.retrieve( + retrieval_method=retrieval_model["search_method"], + dataset_id=dataset.id, + query=query, + top_k=retrieval_model.get("top_k") or 2, + score_threshold=retrieval_model.get("score_threshold", 0.0) + if retrieval_model["score_threshold_enabled"] + else 0.0, + reranking_model=retrieval_model.get("reranking_model", None) + if retrieval_model["reranking_enable"] + else None, + reranking_mode=retrieval_model.get("reranking_mode") or "reranking_model", + weights=retrieval_model.get("weights", None), + ) + + all_documents.extend(documents) + + def to_dataset_retriever_tool( + self, + tenant_id: str, + dataset_ids: list[str], + retrieve_config: DatasetRetrieveConfigEntity, + return_resource: bool, + invoke_from: InvokeFrom, + hit_callback: DatasetIndexToolCallbackHandler, + ) -> Optional[list[DatasetRetrieverBaseTool]]: + """ + A dataset tool is a tool that can be used to retrieve information from a dataset + :param tenant_id: tenant id + :param dataset_ids: dataset ids + :param retrieve_config: retrieve config + :param return_resource: return resource + :param invoke_from: invoke from + :param hit_callback: hit callback + """ + tools = [] + available_datasets = [] + for dataset_id in dataset_ids: + # get dataset from dataset id + dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() + + # pass if dataset is not available + if not dataset: + continue + + # pass if dataset is not available + if dataset and dataset.provider != "external" and dataset.available_document_count == 0: + continue + + available_datasets.append(dataset) + + if retrieve_config.retrieve_strategy == DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE: + # get retrieval model config + default_retrieval_model = { + "search_method": RetrievalMethod.SEMANTIC_SEARCH.value, + "reranking_enable": False, + "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""}, + "top_k": 2, + "score_threshold_enabled": False, + } + + for dataset in available_datasets: + retrieval_model_config = dataset.retrieval_model or default_retrieval_model + + # get top k + top_k = retrieval_model_config["top_k"] + + # get score threshold + score_threshold = None + score_threshold_enabled = retrieval_model_config.get("score_threshold_enabled") + if score_threshold_enabled: + score_threshold = retrieval_model_config.get("score_threshold") + + tool = DatasetRetrieverTool.from_dataset( + dataset=dataset, + top_k=top_k, + score_threshold=score_threshold, + hit_callbacks=[hit_callback], + return_resource=return_resource, + retriever_from=invoke_from.to_source(), + ) + + tools.append(tool) + elif retrieve_config.retrieve_strategy == DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE: + if retrieve_config.reranking_model is not None: + tool = DatasetMultiRetrieverTool.from_dataset( + dataset_ids=[dataset.id for dataset in available_datasets], + tenant_id=tenant_id, + top_k=retrieve_config.top_k or 2, + score_threshold=retrieve_config.score_threshold, + hit_callbacks=[hit_callback], + return_resource=return_resource, + retriever_from=invoke_from.to_source(), + reranking_provider_name=retrieve_config.reranking_model.get("reranking_provider_name"), + reranking_model_name=retrieve_config.reranking_model.get("reranking_model_name"), + ) + + tools.append(tool) + + return tools + + def calculate_keyword_score(self, query: str, documents: list[Document], top_k: int) -> list[Document]: + """ + Calculate keywords scores + :param query: search query + :param documents: documents for reranking + + :return: + """ + keyword_table_handler = JiebaKeywordTableHandler() + query_keywords = keyword_table_handler.extract_keywords(query, None) + documents_keywords = [] + for document in documents: + if document.metadata is not None: + # get the document keywords + document_keywords = keyword_table_handler.extract_keywords(document.page_content, None) + document.metadata["keywords"] = document_keywords + documents_keywords.append(document_keywords) + + # Counter query keywords(TF) + query_keyword_counts = Counter(query_keywords) + + # total documents + total_documents = len(documents) + + # calculate all documents' keywords IDF + all_keywords = set() + for document_keywords in documents_keywords: + all_keywords.update(document_keywords) + + keyword_idf = {} + for keyword in all_keywords: + # calculate include query keywords' documents + doc_count_containing_keyword = sum(1 for doc_keywords in documents_keywords if keyword in doc_keywords) + # IDF + keyword_idf[keyword] = math.log((1 + total_documents) / (1 + doc_count_containing_keyword)) + 1 + + query_tfidf = {} + + for keyword, count in query_keyword_counts.items(): + tf = count + idf = keyword_idf.get(keyword, 0) + query_tfidf[keyword] = tf * idf + + # calculate all documents' TF-IDF + documents_tfidf = [] + for document_keywords in documents_keywords: + document_keyword_counts = Counter(document_keywords) + document_tfidf = {} + for keyword, count in document_keyword_counts.items(): + tf = count + idf = keyword_idf.get(keyword, 0) + document_tfidf[keyword] = tf * idf + documents_tfidf.append(document_tfidf) + + def cosine_similarity(vec1, vec2): + intersection = set(vec1.keys()) & set(vec2.keys()) + numerator = sum(vec1[x] * vec2[x] for x in intersection) + + sum1 = sum(vec1[x] ** 2 for x in vec1) + sum2 = sum(vec2[x] ** 2 for x in vec2) + denominator = math.sqrt(sum1) * math.sqrt(sum2) + + if not denominator: + return 0.0 + else: + return float(numerator) / denominator + + similarities = [] + for document_tfidf in documents_tfidf: + similarity = cosine_similarity(query_tfidf, document_tfidf) + similarities.append(similarity) + + for document, score in zip(documents, similarities): + # format document + if document.metadata is not None: + document.metadata["score"] = score + documents = sorted(documents, key=lambda x: x.metadata.get("score", 0) if x.metadata else 0, reverse=True) + return documents[:top_k] if top_k else documents + + def calculate_vector_score( + self, all_documents: list[Document], top_k: int, score_threshold: float + ) -> list[Document]: + filter_documents = [] + for document in all_documents: + if score_threshold is None or (document.metadata and document.metadata.get("score", 0) >= score_threshold): + filter_documents.append(document) + + if not filter_documents: + return [] + filter_documents = sorted( + filter_documents, key=lambda x: x.metadata.get("score", 0) if x.metadata else 0, reverse=True + ) + return filter_documents[:top_k] if top_k else filter_documents diff --git a/api/core/rag/retrieval/output_parser/__init__.py b/api/core/rag/retrieval/output_parser/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/retrieval/output_parser/react_output.py b/api/core/rag/retrieval/output_parser/react_output.py new file mode 100644 index 0000000000000000000000000000000000000000..9a14d417164e624b44059d5a7c1624c0773b4ad5 --- /dev/null +++ b/api/core/rag/retrieval/output_parser/react_output.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import NamedTuple, Union + + +@dataclass +class ReactAction: + """A full description of an action for an ReactAction to execute.""" + + tool: str + """The name of the Tool to execute.""" + tool_input: Union[str, dict] + """The input to pass in to the Tool.""" + log: str + """Additional information to log about the action.""" + + +class ReactFinish(NamedTuple): + """The final return value of an ReactFinish.""" + + return_values: dict + """Dictionary of return values.""" + log: str + """Additional information to log about the return value""" diff --git a/api/core/rag/retrieval/output_parser/structured_chat.py b/api/core/rag/retrieval/output_parser/structured_chat.py new file mode 100644 index 0000000000000000000000000000000000000000..7fc78bce8357da1656463678dd5a633a832a2e60 --- /dev/null +++ b/api/core/rag/retrieval/output_parser/structured_chat.py @@ -0,0 +1,23 @@ +import json +import re +from typing import Union + +from core.rag.retrieval.output_parser.react_output import ReactAction, ReactFinish + + +class StructuredChatOutputParser: + def parse(self, text: str) -> Union[ReactAction, ReactFinish]: + try: + action_match = re.search(r"```(\w*)\n?({.*?)```", text, re.DOTALL) + if action_match is not None: + response = json.loads(action_match.group(2).strip(), strict=False) + if isinstance(response, list): + response = response[0] + if response["action"] == "Final Answer": + return ReactFinish({"output": response["action_input"]}, text) + else: + return ReactAction(response["action"], response.get("action_input", {}), text) + else: + return ReactFinish({"output": text}, text) + except Exception as e: + raise ValueError(f"Could not parse LLM output: {text}") diff --git a/api/core/rag/retrieval/retrieval_methods.py b/api/core/rag/retrieval/retrieval_methods.py new file mode 100644 index 0000000000000000000000000000000000000000..eaa00bca884a7c03e615b173a2d9d2d06d83c4e6 --- /dev/null +++ b/api/core/rag/retrieval/retrieval_methods.py @@ -0,0 +1,15 @@ +from enum import Enum + + +class RetrievalMethod(Enum): + SEMANTIC_SEARCH = "semantic_search" + FULL_TEXT_SEARCH = "full_text_search" + HYBRID_SEARCH = "hybrid_search" + + @staticmethod + def is_support_semantic_search(retrieval_method: str) -> bool: + return retrieval_method in {RetrievalMethod.SEMANTIC_SEARCH.value, RetrievalMethod.HYBRID_SEARCH.value} + + @staticmethod + def is_support_fulltext_search(retrieval_method: str) -> bool: + return retrieval_method in {RetrievalMethod.FULL_TEXT_SEARCH.value, RetrievalMethod.HYBRID_SEARCH.value} diff --git a/api/core/rag/retrieval/router/multi_dataset_function_call_router.py b/api/core/rag/retrieval/router/multi_dataset_function_call_router.py new file mode 100644 index 0000000000000000000000000000000000000000..b008d0df9c2f0e33b4bdc903fe0637dcb6cc024f --- /dev/null +++ b/api/core/rag/retrieval/router/multi_dataset_function_call_router.py @@ -0,0 +1,45 @@ +from typing import Union, cast + +from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity +from core.model_manager import ModelInstance +from core.model_runtime.entities.llm_entities import LLMResult +from core.model_runtime.entities.message_entities import PromptMessageTool, SystemPromptMessage, UserPromptMessage + + +class FunctionCallMultiDatasetRouter: + def invoke( + self, + query: str, + dataset_tools: list[PromptMessageTool], + model_config: ModelConfigWithCredentialsEntity, + model_instance: ModelInstance, + ) -> Union[str, None]: + """Given input, decided what to do. + Returns: + Action specifying what tool to use. + """ + if len(dataset_tools) == 0: + return None + elif len(dataset_tools) == 1: + return dataset_tools[0].name + + try: + prompt_messages = [ + SystemPromptMessage(content="You are a helpful AI assistant."), + UserPromptMessage(content=query), + ] + result = cast( + LLMResult, + model_instance.invoke_llm( + prompt_messages=prompt_messages, + tools=dataset_tools, + stream=False, + model_parameters={"temperature": 0.2, "top_p": 0.3, "max_tokens": 1500}, + ), + ) + if result.message.tool_calls: + # get retrieval model config + return result.message.tool_calls[0].function.name + return None + except Exception as e: + return None diff --git a/api/core/rag/retrieval/router/multi_dataset_react_route.py b/api/core/rag/retrieval/router/multi_dataset_react_route.py new file mode 100644 index 0000000000000000000000000000000000000000..05e8d043dfe741b4d08611bf7a3cc6aa5c9018e8 --- /dev/null +++ b/api/core/rag/retrieval/router/multi_dataset_react_route.py @@ -0,0 +1,250 @@ +from collections.abc import Generator, Sequence +from typing import Union, cast + +from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity +from core.model_manager import ModelInstance +from core.model_runtime.entities.llm_entities import LLMResult, LLMUsage +from core.model_runtime.entities.message_entities import PromptMessage, PromptMessageRole, PromptMessageTool +from core.prompt.advanced_prompt_transform import AdvancedPromptTransform +from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate +from core.rag.retrieval.output_parser.react_output import ReactAction +from core.rag.retrieval.output_parser.structured_chat import StructuredChatOutputParser +from core.workflow.nodes.llm import LLMNode + +PREFIX = """Respond to the human as helpfully and accurately as possible. You have access to the following tools:""" + +SUFFIX = """Begin! Reminder to ALWAYS respond with a valid json blob of a single action. Use tools if necessary. Respond directly if appropriate. Format is Action:```$JSON_BLOB```then Observation:. +Thought:""" # noqa: E501 + +FORMAT_INSTRUCTIONS = """Use a json blob to specify a tool by providing an action key (tool name) and an action_input key (tool input). +The nouns in the format of "Thought", "Action", "Action Input", "Final Answer" must be expressed in English. +Valid "action" values: "Final Answer" or {tool_names} + +Provide only ONE action per $JSON_BLOB, as shown: + +``` +{{ + "action": $TOOL_NAME, + "action_input": $INPUT +}} +``` + +Follow this format: + +Question: input question to answer +Thought: consider previous and subsequent steps +Action: +``` +$JSON_BLOB +``` +Observation: action result +... (repeat Thought/Action/Observation N times) +Thought: I know what to respond +Action: +``` +{{ + "action": "Final Answer", + "action_input": "Final response to human" +}} +```""" # noqa: E501 + + +class ReactMultiDatasetRouter: + def invoke( + self, + query: str, + dataset_tools: list[PromptMessageTool], + model_config: ModelConfigWithCredentialsEntity, + model_instance: ModelInstance, + user_id: str, + tenant_id: str, + ) -> Union[str, None]: + """Given input, decided what to do. + Returns: + Action specifying what tool to use. + """ + if len(dataset_tools) == 0: + return None + elif len(dataset_tools) == 1: + return dataset_tools[0].name + + try: + return self._react_invoke( + query=query, + model_config=model_config, + model_instance=model_instance, + tools=dataset_tools, + user_id=user_id, + tenant_id=tenant_id, + ) + except Exception as e: + return None + + def _react_invoke( + self, + query: str, + model_config: ModelConfigWithCredentialsEntity, + model_instance: ModelInstance, + tools: Sequence[PromptMessageTool], + user_id: str, + tenant_id: str, + prefix: str = PREFIX, + suffix: str = SUFFIX, + format_instructions: str = FORMAT_INSTRUCTIONS, + ) -> Union[str, None]: + prompt: Union[list[ChatModelMessage], CompletionModelPromptTemplate] + if model_config.mode == "chat": + prompt = self.create_chat_prompt( + query=query, + tools=tools, + prefix=prefix, + suffix=suffix, + format_instructions=format_instructions, + ) + else: + prompt = self.create_completion_prompt( + tools=tools, + prefix=prefix, + format_instructions=format_instructions, + ) + stop = ["Observation:"] + # handle invoke result + prompt_transform = AdvancedPromptTransform() + prompt_messages = prompt_transform.get_prompt( + prompt_template=prompt, + inputs={}, + query="", + files=[], + context="", + memory_config=None, + memory=None, + model_config=model_config, + ) + result_text, usage = self._invoke_llm( + completion_param=model_config.parameters, + model_instance=model_instance, + prompt_messages=prompt_messages, + stop=stop, + user_id=user_id, + tenant_id=tenant_id, + ) + output_parser = StructuredChatOutputParser() + react_decision = output_parser.parse(result_text) + if isinstance(react_decision, ReactAction): + return react_decision.tool + return None + + def _invoke_llm( + self, + completion_param: dict, + model_instance: ModelInstance, + prompt_messages: list[PromptMessage], + stop: list[str], + user_id: str, + tenant_id: str, + ) -> tuple[str, LLMUsage]: + """ + Invoke large language model + :param model_instance: model instance + :param prompt_messages: prompt messages + :param stop: stop + :return: + """ + invoke_result = cast( + Generator[LLMResult, None, None], + model_instance.invoke_llm( + prompt_messages=prompt_messages, + model_parameters=completion_param, + stop=stop, + stream=True, + user=user_id, + ), + ) + + # handle invoke result + text, usage = self._handle_invoke_result(invoke_result=invoke_result) + + # deduct quota + LLMNode.deduct_llm_quota(tenant_id=tenant_id, model_instance=model_instance, usage=usage) + + return text, usage + + def _handle_invoke_result(self, invoke_result: Generator) -> tuple[str, LLMUsage]: + """ + Handle invoke result + :param invoke_result: invoke result + :return: + """ + model = None + prompt_messages: list[PromptMessage] = [] + full_text = "" + usage = None + for result in invoke_result: + text = result.delta.message.content + full_text += text + + if not model: + model = result.model + + if not prompt_messages: + prompt_messages = result.prompt_messages + + if not usage and result.delta.usage: + usage = result.delta.usage + + if not usage: + usage = LLMUsage.empty_usage() + + return full_text, usage + + def create_chat_prompt( + self, + query: str, + tools: Sequence[PromptMessageTool], + prefix: str = PREFIX, + suffix: str = SUFFIX, + format_instructions: str = FORMAT_INSTRUCTIONS, + ) -> list[ChatModelMessage]: + tool_strings = [] + for tool in tools: + tool_strings.append( + f"{tool.name}: {tool.description}, args: {{'query': {{'title': 'Query'," + f" 'description': 'Query for the dataset to be used to retrieve the dataset.', 'type': 'string'}}}}" + ) + formatted_tools = "\n".join(tool_strings) + unique_tool_names = {tool.name for tool in tools} + tool_names = ", ".join('"' + name + '"' for name in unique_tool_names) + format_instructions = format_instructions.format(tool_names=tool_names) + template = "\n\n".join([prefix, formatted_tools, format_instructions, suffix]) + prompt_messages = [] + system_prompt_messages = ChatModelMessage(role=PromptMessageRole.SYSTEM, text=template) + prompt_messages.append(system_prompt_messages) + user_prompt_message = ChatModelMessage(role=PromptMessageRole.USER, text=query) + prompt_messages.append(user_prompt_message) + return prompt_messages + + def create_completion_prompt( + self, + tools: Sequence[PromptMessageTool], + prefix: str = PREFIX, + format_instructions: str = FORMAT_INSTRUCTIONS, + ) -> CompletionModelPromptTemplate: + """Create prompt in the style of the zero shot agent. + + Args: + tools: List of tools the agent will have access to, used to format the + prompt. + prefix: String to put before the list of tools. + Returns: + A PromptTemplate with the template assembled from the pieces here. + """ + suffix = """Begin! Reminder to ALWAYS respond with a valid json blob of a single action. Use tools if necessary. Respond directly if appropriate. Format is Action:```$JSON_BLOB```then Observation:. +Question: {input} +Thought: {agent_scratchpad} +""" # noqa: E501 + + tool_strings = "\n".join([f"{tool.name}: {tool.description}" for tool in tools]) + tool_names = ", ".join([tool.name for tool in tools]) + format_instructions = format_instructions.format(tool_names=tool_names) + template = "\n\n".join([prefix, tool_strings, format_instructions, suffix]) + return CompletionModelPromptTemplate(text=template) diff --git a/api/core/rag/splitter/__init__.py b/api/core/rag/splitter/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/rag/splitter/fixed_text_splitter.py b/api/core/rag/splitter/fixed_text_splitter.py new file mode 100644 index 0000000000000000000000000000000000000000..3376bd7f75dd96ea8a836c7b05d854d6649a86ce --- /dev/null +++ b/api/core/rag/splitter/fixed_text_splitter.py @@ -0,0 +1,112 @@ +"""Functionality for splitting text.""" + +from __future__ import annotations + +from typing import Any, Optional + +from core.model_manager import ModelInstance +from core.model_runtime.model_providers.__base.tokenizers.gpt2_tokenzier import GPT2Tokenizer +from core.rag.splitter.text_splitter import ( + TS, + Collection, + Literal, + RecursiveCharacterTextSplitter, + Set, + TokenTextSplitter, + Union, +) + + +class EnhanceRecursiveCharacterTextSplitter(RecursiveCharacterTextSplitter): + """ + This class is used to implement from_gpt2_encoder, to prevent using of tiktoken + """ + + @classmethod + def from_encoder( + cls: type[TS], + embedding_model_instance: Optional[ModelInstance], + allowed_special: Union[Literal["all"], Set[str]] = set(), # noqa: UP037 + disallowed_special: Union[Literal["all"], Collection[str]] = "all", # noqa: UP037 + **kwargs: Any, + ): + def _token_encoder(text: str) -> int: + if not text: + return 0 + + if embedding_model_instance: + return embedding_model_instance.get_text_embedding_num_tokens(texts=[text]) + else: + return GPT2Tokenizer.get_num_tokens(text) + + if issubclass(cls, TokenTextSplitter): + extra_kwargs = { + "model_name": embedding_model_instance.model if embedding_model_instance else "gpt2", + "allowed_special": allowed_special, + "disallowed_special": disallowed_special, + } + kwargs = {**kwargs, **extra_kwargs} + + return cls(length_function=_token_encoder, **kwargs) + + +class FixedRecursiveCharacterTextSplitter(EnhanceRecursiveCharacterTextSplitter): + def __init__(self, fixed_separator: str = "\n\n", separators: Optional[list[str]] = None, **kwargs: Any): + """Create a new TextSplitter.""" + super().__init__(**kwargs) + self._fixed_separator = fixed_separator + self._separators = separators or ["\n\n", "\n", " ", ""] + + def split_text(self, text: str) -> list[str]: + """Split incoming text and return chunks.""" + if self._fixed_separator: + chunks = text.split(self._fixed_separator) + else: + chunks = [text] + + final_chunks = [] + for chunk in chunks: + if self._length_function(chunk) > self._chunk_size: + final_chunks.extend(self.recursive_split_text(chunk)) + else: + final_chunks.append(chunk) + + return final_chunks + + def recursive_split_text(self, text: str) -> list[str]: + """Split incoming text and return chunks.""" + final_chunks = [] + # Get appropriate separator to use + separator = self._separators[-1] + for _s in self._separators: + if _s == "": + separator = _s + break + if _s in text: + separator = _s + break + # Now that we have the separator, split the text + if separator: + splits = text.split(separator) + else: + splits = list(text) + # Now go merging things, recursively splitting longer texts. + _good_splits = [] + _good_splits_lengths = [] # cache the lengths of the splits + for s in splits: + s_len = self._length_function(s) + if s_len < self._chunk_size: + _good_splits.append(s) + _good_splits_lengths.append(s_len) + else: + if _good_splits: + merged_text = self._merge_splits(_good_splits, separator, _good_splits_lengths) + final_chunks.extend(merged_text) + _good_splits = [] + _good_splits_lengths = [] + other_info = self.recursive_split_text(s) + final_chunks.extend(other_info) + if _good_splits: + merged_text = self._merge_splits(_good_splits, separator, _good_splits_lengths) + final_chunks.extend(merged_text) + return final_chunks diff --git a/api/core/rag/splitter/text_splitter.py b/api/core/rag/splitter/text_splitter.py new file mode 100644 index 0000000000000000000000000000000000000000..4bfa541fd454ad747bc429778ef66139441ed3e7 --- /dev/null +++ b/api/core/rag/splitter/text_splitter.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import copy +import logging +import re +from abc import ABC, abstractmethod +from collections.abc import Callable, Collection, Iterable, Sequence, Set +from dataclasses import dataclass +from typing import ( + Any, + Literal, + Optional, + TypedDict, + TypeVar, + Union, +) + +from core.rag.models.document import BaseDocumentTransformer, Document + +logger = logging.getLogger(__name__) + +TS = TypeVar("TS", bound="TextSplitter") + + +def _split_text_with_regex(text: str, separator: str, keep_separator: bool) -> list[str]: + # Now that we have the separator, split the text + if separator: + if keep_separator: + # The parentheses in the pattern keep the delimiters in the result. + _splits = re.split(f"({re.escape(separator)})", text) + splits = [_splits[i - 1] + _splits[i] for i in range(1, len(_splits), 2)] + if len(_splits) % 2 != 0: + splits += _splits[-1:] + else: + splits = re.split(separator, text) + else: + splits = list(text) + return [s for s in splits if (s not in {"", "\n"})] + + +class TextSplitter(BaseDocumentTransformer, ABC): + """Interface for splitting text into chunks.""" + + def __init__( + self, + chunk_size: int = 4000, + chunk_overlap: int = 200, + length_function: Callable[[str], int] = len, + keep_separator: bool = False, + add_start_index: bool = False, + ) -> None: + """Create a new TextSplitter. + + Args: + chunk_size: Maximum size of chunks to return + chunk_overlap: Overlap in characters between chunks + length_function: Function that measures the length of given chunks + keep_separator: Whether to keep the separator in the chunks + add_start_index: If `True`, includes chunk's start index in metadata + """ + if chunk_overlap > chunk_size: + raise ValueError( + f"Got a larger chunk overlap ({chunk_overlap}) than chunk size ({chunk_size}), should be smaller." + ) + self._chunk_size = chunk_size + self._chunk_overlap = chunk_overlap + self._length_function = length_function + self._keep_separator = keep_separator + self._add_start_index = add_start_index + + @abstractmethod + def split_text(self, text: str) -> list[str]: + """Split text into multiple components.""" + + def create_documents(self, texts: list[str], metadatas: Optional[list[dict]] = None) -> list[Document]: + """Create documents from a list of texts.""" + _metadatas = metadatas or [{}] * len(texts) + documents = [] + for i, text in enumerate(texts): + index = -1 + for chunk in self.split_text(text): + metadata = copy.deepcopy(_metadatas[i]) + if self._add_start_index: + index = text.find(chunk, index + 1) + metadata["start_index"] = index + new_doc = Document(page_content=chunk, metadata=metadata) + documents.append(new_doc) + return documents + + def split_documents(self, documents: Iterable[Document]) -> list[Document]: + """Split documents.""" + texts, metadatas = [], [] + for doc in documents: + texts.append(doc.page_content) + metadatas.append(doc.metadata or {}) + return self.create_documents(texts, metadatas=metadatas) + + def _join_docs(self, docs: list[str], separator: str) -> Optional[str]: + text = separator.join(docs) + text = text.strip() + if text == "": + return None + else: + return text + + def _merge_splits(self, splits: Iterable[str], separator: str, lengths: list[int]) -> list[str]: + # We now want to combine these smaller pieces into medium size + # chunks to send to the LLM. + separator_len = self._length_function(separator) + + docs = [] + current_doc: list[str] = [] + total = 0 + index = 0 + for d in splits: + _len = lengths[index] + if total + _len + (separator_len if len(current_doc) > 0 else 0) > self._chunk_size: + if total > self._chunk_size: + logger.warning( + f"Created a chunk of size {total}, which is longer than the specified {self._chunk_size}" + ) + if len(current_doc) > 0: + doc = self._join_docs(current_doc, separator) + if doc is not None: + docs.append(doc) + # Keep on popping if: + # - we have a larger chunk than in the chunk overlap + # - or if we still have any chunks and the length is long + while total > self._chunk_overlap or ( + total + _len + (separator_len if len(current_doc) > 0 else 0) > self._chunk_size and total > 0 + ): + total -= self._length_function(current_doc[0]) + (separator_len if len(current_doc) > 1 else 0) + current_doc = current_doc[1:] + current_doc.append(d) + total += _len + (separator_len if len(current_doc) > 1 else 0) + index += 1 + doc = self._join_docs(current_doc, separator) + if doc is not None: + docs.append(doc) + return docs + + @classmethod + def from_huggingface_tokenizer(cls, tokenizer: Any, **kwargs: Any) -> TextSplitter: + """Text splitter that uses HuggingFace tokenizer to count length.""" + try: + from transformers import PreTrainedTokenizerBase # type: ignore + + if not isinstance(tokenizer, PreTrainedTokenizerBase): + raise ValueError("Tokenizer received was not an instance of PreTrainedTokenizerBase") + + def _huggingface_tokenizer_length(text: str) -> int: + return len(tokenizer.encode(text)) + + except ImportError: + raise ValueError( + "Could not import transformers python package. Please install it with `pip install transformers`." + ) + return cls(length_function=_huggingface_tokenizer_length, **kwargs) + + @classmethod + def from_tiktoken_encoder( + cls: type[TS], + encoding_name: str = "gpt2", + model_name: Optional[str] = None, + allowed_special: Union[Literal["all"], Set[str]] = set(), + disallowed_special: Union[Literal["all"], Collection[str]] = "all", + **kwargs: Any, + ) -> TS: + """Text splitter that uses tiktoken encoder to count length.""" + try: + import tiktoken + except ImportError: + raise ImportError( + "Could not import tiktoken python package. " + "This is needed in order to calculate max_tokens_for_prompt. " + "Please install it with `pip install tiktoken`." + ) + + if model_name is not None: + enc = tiktoken.encoding_for_model(model_name) + else: + enc = tiktoken.get_encoding(encoding_name) + + def _tiktoken_encoder(text: str) -> int: + return len( + enc.encode( + text, + allowed_special=allowed_special, + disallowed_special=disallowed_special, + ) + ) + + if issubclass(cls, TokenTextSplitter): + extra_kwargs = { + "encoding_name": encoding_name, + "model_name": model_name, + "allowed_special": allowed_special, + "disallowed_special": disallowed_special, + } + kwargs = {**kwargs, **extra_kwargs} + + return cls(length_function=_tiktoken_encoder, **kwargs) + + def transform_documents(self, documents: Sequence[Document], **kwargs: Any) -> Sequence[Document]: + """Transform sequence of documents by splitting them.""" + return self.split_documents(list(documents)) + + async def atransform_documents(self, documents: Sequence[Document], **kwargs: Any) -> Sequence[Document]: + """Asynchronously transform a sequence of documents by splitting them.""" + raise NotImplementedError + + +class CharacterTextSplitter(TextSplitter): + """Splitting text that looks at characters.""" + + def __init__(self, separator: str = "\n\n", **kwargs: Any) -> None: + """Create a new TextSplitter.""" + super().__init__(**kwargs) + self._separator = separator + + def split_text(self, text: str) -> list[str]: + """Split incoming text and return chunks.""" + # First we naively split the large input into a bunch of smaller ones. + splits = _split_text_with_regex(text, self._separator, self._keep_separator) + _separator = "" if self._keep_separator else self._separator + _good_splits_lengths = [] # cache the lengths of the splits + for split in splits: + _good_splits_lengths.append(self._length_function(split)) + return self._merge_splits(splits, _separator, _good_splits_lengths) + + +class LineType(TypedDict): + """Line type as typed dict.""" + + metadata: dict[str, str] + content: str + + +class HeaderType(TypedDict): + """Header type as typed dict.""" + + level: int + name: str + data: str + + +class MarkdownHeaderTextSplitter: + """Splitting markdown files based on specified headers.""" + + def __init__(self, headers_to_split_on: list[tuple[str, str]], return_each_line: bool = False): + """Create a new MarkdownHeaderTextSplitter. + + Args: + headers_to_split_on: Headers we want to track + return_each_line: Return each line w/ associated headers + """ + # Output line-by-line or aggregated into chunks w/ common headers + self.return_each_line = return_each_line + # Given the headers we want to split on, + # (e.g., "#, ##, etc") order by length + self.headers_to_split_on = sorted(headers_to_split_on, key=lambda split: len(split[0]), reverse=True) + + def aggregate_lines_to_chunks(self, lines: list[LineType]) -> list[Document]: + """Combine lines with common metadata into chunks + Args: + lines: Line of text / associated header metadata + """ + aggregated_chunks: list[LineType] = [] + + for line in lines: + if aggregated_chunks and aggregated_chunks[-1]["metadata"] == line["metadata"]: + # If the last line in the aggregated list + # has the same metadata as the current line, + # append the current content to the last lines's content + aggregated_chunks[-1]["content"] += " \n" + line["content"] + else: + # Otherwise, append the current line to the aggregated list + aggregated_chunks.append(line) + + return [Document(page_content=chunk["content"], metadata=chunk["metadata"]) for chunk in aggregated_chunks] + + def split_text(self, text: str) -> list[Document]: + """Split markdown file + Args: + text: Markdown file""" + + # Split the input text by newline character ("\n"). + lines = text.split("\n") + # Final output + lines_with_metadata: list[LineType] = [] + # Content and metadata of the chunk currently being processed + current_content: list[str] = [] + current_metadata: dict[str, str] = {} + # Keep track of the nested header structure + # header_stack: List[Dict[str, Union[int, str]]] = [] + header_stack: list[HeaderType] = [] + initial_metadata: dict[str, str] = {} + + for line in lines: + stripped_line = line.strip() + # Check each line against each of the header types (e.g., #, ##) + for sep, name in self.headers_to_split_on: + # Check if line starts with a header that we intend to split on + if stripped_line.startswith(sep) and ( + # Header with no text OR header is followed by space + # Both are valid conditions that sep is being used a header + len(stripped_line) == len(sep) or stripped_line[len(sep)] == " " + ): + # Ensure we are tracking the header as metadata + if name is not None: + # Get the current header level + current_header_level = sep.count("#") + + # Pop out headers of lower or same level from the stack + while header_stack and header_stack[-1]["level"] >= current_header_level: + # We have encountered a new header + # at the same or higher level + popped_header = header_stack.pop() + # Clear the metadata for the + # popped header in initial_metadata + if popped_header["name"] in initial_metadata: + initial_metadata.pop(popped_header["name"]) + + # Push the current header to the stack + header: HeaderType = { + "level": current_header_level, + "name": name, + "data": stripped_line[len(sep) :].strip(), + } + header_stack.append(header) + # Update initial_metadata with the current header + initial_metadata[name] = header["data"] + + # Add the previous line to the lines_with_metadata + # only if current_content is not empty + if current_content: + lines_with_metadata.append( + { + "content": "\n".join(current_content), + "metadata": current_metadata.copy(), + } + ) + current_content.clear() + + break + else: + if stripped_line: + current_content.append(stripped_line) + elif current_content: + lines_with_metadata.append( + { + "content": "\n".join(current_content), + "metadata": current_metadata.copy(), + } + ) + current_content.clear() + + current_metadata = initial_metadata.copy() + + if current_content: + lines_with_metadata.append({"content": "\n".join(current_content), "metadata": current_metadata}) + + # lines_with_metadata has each line with associated header metadata + # aggregate these into chunks based on common metadata + if not self.return_each_line: + return self.aggregate_lines_to_chunks(lines_with_metadata) + else: + return [ + Document(page_content=chunk["content"], metadata=chunk["metadata"]) for chunk in lines_with_metadata + ] + + +# should be in newer Python versions (3.10+) +# @dataclass(frozen=True, kw_only=True, slots=True) +@dataclass(frozen=True) +class Tokenizer: + chunk_overlap: int + tokens_per_chunk: int + decode: Callable[[list[int]], str] + encode: Callable[[str], list[int]] + + +def split_text_on_tokens(*, text: str, tokenizer: Tokenizer) -> list[str]: + """Split incoming text and return chunks using tokenizer.""" + splits: list[str] = [] + input_ids = tokenizer.encode(text) + start_idx = 0 + cur_idx = min(start_idx + tokenizer.tokens_per_chunk, len(input_ids)) + chunk_ids = input_ids[start_idx:cur_idx] + while start_idx < len(input_ids): + splits.append(tokenizer.decode(chunk_ids)) + start_idx += tokenizer.tokens_per_chunk - tokenizer.chunk_overlap + cur_idx = min(start_idx + tokenizer.tokens_per_chunk, len(input_ids)) + chunk_ids = input_ids[start_idx:cur_idx] + return splits + + +class TokenTextSplitter(TextSplitter): + """Splitting text to tokens using model tokenizer.""" + + def __init__( + self, + encoding_name: str = "gpt2", + model_name: Optional[str] = None, + allowed_special: Union[Literal["all"], Set[str]] = set(), + disallowed_special: Union[Literal["all"], Collection[str]] = "all", + **kwargs: Any, + ) -> None: + """Create a new TextSplitter.""" + super().__init__(**kwargs) + try: + import tiktoken + except ImportError: + raise ImportError( + "Could not import tiktoken python package. " + "This is needed in order to for TokenTextSplitter. " + "Please install it with `pip install tiktoken`." + ) + + if model_name is not None: + enc = tiktoken.encoding_for_model(model_name) + else: + enc = tiktoken.get_encoding(encoding_name) + self._tokenizer = enc + self._allowed_special = allowed_special + self._disallowed_special = disallowed_special + + def split_text(self, text: str) -> list[str]: + def _encode(_text: str) -> list[int]: + return self._tokenizer.encode( + _text, + allowed_special=self._allowed_special, + disallowed_special=self._disallowed_special, + ) + + tokenizer = Tokenizer( + chunk_overlap=self._chunk_overlap, + tokens_per_chunk=self._chunk_size, + decode=self._tokenizer.decode, + encode=_encode, + ) + + return split_text_on_tokens(text=text, tokenizer=tokenizer) + + +class RecursiveCharacterTextSplitter(TextSplitter): + """Splitting text by recursively look at characters. + + Recursively tries to split by different characters to find one + that works. + """ + + def __init__( + self, + separators: Optional[list[str]] = None, + keep_separator: bool = True, + **kwargs: Any, + ) -> None: + """Create a new TextSplitter.""" + super().__init__(keep_separator=keep_separator, **kwargs) + self._separators = separators or ["\n\n", "\n", " ", ""] + + def _split_text(self, text: str, separators: list[str]) -> list[str]: + final_chunks = [] + separator = separators[-1] + new_separators = [] + + for i, _s in enumerate(separators): + if _s == "": + separator = _s + break + if re.search(_s, text): + separator = _s + new_separators = separators[i + 1 :] + break + + splits = _split_text_with_regex(text, separator, self._keep_separator) + _good_splits = [] + _good_splits_lengths = [] # cache the lengths of the splits + _separator = "" if self._keep_separator else separator + + for s in splits: + s_len = self._length_function(s) + if s_len < self._chunk_size: + _good_splits.append(s) + _good_splits_lengths.append(s_len) + else: + if _good_splits: + merged_text = self._merge_splits(_good_splits, _separator, _good_splits_lengths) + final_chunks.extend(merged_text) + _good_splits = [] + _good_splits_lengths = [] + if not new_separators: + final_chunks.append(s) + else: + other_info = self._split_text(s, new_separators) + final_chunks.extend(other_info) + + if _good_splits: + merged_text = self._merge_splits(_good_splits, _separator, _good_splits_lengths) + final_chunks.extend(merged_text) + + return final_chunks + + def split_text(self, text: str) -> list[str]: + return self._split_text(text, self._separators) diff --git a/api/core/tools/README.md b/api/core/tools/README.md new file mode 100644 index 0000000000000000000000000000000000000000..b5d0a30d348a9d10811d3303ce58b78b4c8827e6 --- /dev/null +++ b/api/core/tools/README.md @@ -0,0 +1,25 @@ +# Tools + +This module implements built-in tools used in Agent Assistants and Workflows within Dify. You could define and display your own tools in this module, without modifying the frontend logic. This decoupling allows for easier horizontal scaling of Dify's capabilities. + +## Feature Introduction + +The tools provided for Agents and Workflows are currently divided into two categories: +- `Built-in Tools` are internally implemented within our product and are hardcoded for use in Agents and Workflows. +- `Api-Based Tools` leverage third-party APIs for implementation. You don't need to code to integrate these -- simply provide interface definitions in formats like `OpenAPI` , `Swagger`, or the `OpenAI-plugin` on the front-end. + +### Built-in Tool Providers +![Alt text](docs/images/index/image.png) + +### API Tool Providers +![Alt text](docs/images/index/image-1.png) + +## Tool Integration + +To enable developers to build flexible and powerful tools, we provide two guides: + +### [Quick Integration 👈🏻](./docs/en_US/tool_scale_out.md) +Quick integration aims at quickly getting you up to speed with tool integration by walking over an example Google Search tool. + +### [Advanced Integration 👈🏻](./docs/en_US/advanced_scale_out.md) +Advanced integration will offer a deeper dive into the module interfaces, and explain how to implement more complex capabilities, such as generating images, combining multiple tools, and managing the flow of parameters, images, and files between different tools. \ No newline at end of file diff --git a/api/core/tools/README_CN.md b/api/core/tools/README_CN.md new file mode 100644 index 0000000000000000000000000000000000000000..7e18441131fd75774b06ced425e7ee6b012b819d --- /dev/null +++ b/api/core/tools/README_CN.md @@ -0,0 +1,27 @@ +# Tools + +该模块提供了各Agent和Workflow中会使用的内置工具的调用、鉴权接口,并为 Dify 提供了统一的工具供应商的信息和凭据表单规则。 + +- 一方面将工具和业务代码解耦,方便开发者对模型横向扩展, +- 另一方面提供了只需在后端定义供应商和工具,即可在前端页面直接展示,无需修改前端逻辑。 + +## 功能介绍 + +对于给Agent和Workflow提供的工具,我们当前将其分为两类: +- `Built-in Tools` 内置工具,即Dify内部实现的工具,通过硬编码的方式提供给Agent和Workflow使用。 +- `Api-Based Tools` 基于API的工具,即通过调用第三方API实现的工具,`Api-Based Tool`不需要再额外定义,只需提供`OpenAPI` `Swagger` `OpenAI plugin`等接口文档即可。 + +### 内置工具供应商 +![Alt text](docs/images/index/image.png) + +### API工具供应商 +![Alt text](docs/images/index/image-1.png) + +## 工具接入 +为了实现更灵活更强大的功能,Tools提供了一系列的接口,帮助开发者快速构建想要的工具,本文作为开发者的入门指南,将会以[快速接入](./docs/zh_Hans/tool_scale_out.md)和[高级接入](./docs/zh_Hans/advanced_scale_out.md)两部分介绍如何接入工具。 + +### [快速接入 👈🏻](./docs/zh_Hans/tool_scale_out.md) +快速接入可以帮助你在10~20分钟内完成工具的接入,但是这种接入方式只能实现简单的功能,如果你想要实现更复杂的功能,可以参考下面的高级接入。 + +### [高级接入 👈🏻](./docs/zh_Hans/advanced_scale_out.md) +高级接入将介绍如何实现更复杂的功能配置,包括实现图生图、实现多个工具的组合、实现参数、图片、文件在多个工具之间的流转。 \ No newline at end of file diff --git a/api/core/tools/README_JA.md b/api/core/tools/README_JA.md new file mode 100644 index 0000000000000000000000000000000000000000..39d0bf1762ad0ea31a48ea5573c427b0770877e4 --- /dev/null +++ b/api/core/tools/README_JA.md @@ -0,0 +1,31 @@ +# Tools + +このモジュールは、Difyのエージェントアシスタントやワークフローで使用される組み込みツールを実装しています。このモジュールでは、フロントエンドのロジックを変更することなく、独自のツールを定義し表示することができます。この分離により、Difyの機能を容易に水平方向にスケールアウトできます。 + +## 機能紹介 + +エージェントとワークフロー向けに提供されるツールは、現在2つのカテゴリーに分類されています。 + +- `Built-in Tools`はDify内部で実装され、エージェントとワークフローで使用するためにハードコードされています。 +- `Api-Based Tools`はサードパーティのAPIを利用して実装されています。これらを統合するためのコーディングは不要で、フロントエンドで + `OpenAPI`, `Swagger`または`OpenAI-plugin`などの形式でインターフェース定義を提供するだけです。 + +### 組み込みツールプロバイダー + +![Alt text](docs/images/index/image.png) + +### APIツールプロバイダー + +![Alt text](docs/images/index/image-1.png) + +## ツールの統合 + +開発者が柔軟で強力なツールを構築できるよう、2つのガイドを提供しています。 + +### [クイック統合 👈🏻](./docs/ja_JP/tool_scale_out.md) + +クイック統合は、Google検索ツールの例を通じて、ツール統合の基本をすばやく理解できるようにすることを目的としています。 + +### [高度な統合 👈🏻](./docs/ja_JP/advanced_scale_out.md) + +高度な統合では、モジュールインターフェースについてより深く掘り下げ、画像生成、複数ツールの組み合わせ、異なるツール間でのパラメーター、画像、ファイルのフロー管理など、より複雑な機能の実装方法を説明します。 \ No newline at end of file diff --git a/api/core/tools/__init__.py b/api/core/tools/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/tools/docs/en_US/advanced_scale_out.md b/api/core/tools/docs/en_US/advanced_scale_out.md new file mode 100644 index 0000000000000000000000000000000000000000..644ad291292444d80225b31127e9dc47f492668d --- /dev/null +++ b/api/core/tools/docs/en_US/advanced_scale_out.md @@ -0,0 +1,278 @@ +# Advanced Tool Integration + +Before starting with this advanced guide, please make sure you have a basic understanding of the tool integration process in Dify. Check out [Quick Integration](./tool_scale_out.md) for a quick runthrough. + +## Tool Interface + +We have defined a series of helper methods in the `Tool` class to help developers quickly build more complex tools. + +### Message Return + +Dify supports various message types such as `text`, `link`, `json`, `image`, and `file BLOB`. You can return different types of messages to the LLM and users through the following interfaces. + +Please note, some parameters in the following interfaces will be introduced in later sections. + +#### Image URL +You only need to pass the URL of the image, and Dify will automatically download the image and return it to the user. + +```python + def create_image_message(self, image: str, save_as: str = '') -> ToolInvokeMessage: + """ + create an image message + + :param image: the url of the image + :return: the image message + """ +``` + +#### Link +If you need to return a link, you can use the following interface. + +```python + def create_link_message(self, link: str, save_as: str = '') -> ToolInvokeMessage: + """ + create a link message + + :param link: the url of the link + :return: the link message + """ +``` + +#### Text +If you need to return a text message, you can use the following interface. + +```python + def create_text_message(self, text: str, save_as: str = '') -> ToolInvokeMessage: + """ + create a text message + + :param text: the text of the message + :return: the text message + """ +``` + +#### File BLOB +If you need to return the raw data of a file, such as images, audio, video, PPT, Word, Excel, etc., you can use the following interface. + +- `blob` The raw data of the file, of bytes type +- `meta` The metadata of the file, if you know the type of the file, it is best to pass a `mime_type`, otherwise Dify will use `octet/stream` as the default type + +```python + def create_blob_message(self, blob: bytes, meta: dict = None, save_as: str = '') -> ToolInvokeMessage: + """ + create a blob message + + :param blob: the blob + :return: the blob message + """ +``` + +#### JSON +If you need to return a formatted JSON, you can use the following interface. This is commonly used for data transmission between nodes in a workflow, of course, in agent mode, most LLM are also able to read and understand JSON. + +- `object` A Python dictionary object will be automatically serialized into JSON + +```python + def create_json_message(self, object: dict) -> ToolInvokeMessage: + """ + create a json message + """ +``` + +### Shortcut Tools + +In large model applications, we have two common needs: +- First, summarize a long text in advance, and then pass the summary content to the LLM to prevent the original text from being too long for the LLM to handle +- The content obtained by the tool is a link, and the web page information needs to be crawled before it can be returned to the LLM + +To help developers quickly implement these two needs, we provide the following two shortcut tools. + +#### Text Summary Tool + +This tool takes in an user_id and the text to be summarized, and returns the summarized text. Dify will use the default model of the current workspace to summarize the long text. + +```python + def summary(self, user_id: str, content: str) -> str: + """ + summary the content + + :param user_id: the user id + :param content: the content + :return: the summary + """ +``` + +#### Web Page Crawling Tool + +This tool takes in web page link to be crawled and a user_agent (which can be empty), and returns a string containing the information of the web page. The `user_agent` is an optional parameter that can be used to identify the tool. If not passed, Dify will use the default `user_agent`. + +```python + def get_url(self, url: str, user_agent: str = None) -> str: + """ + get url + """ the crawled result +``` + +### Variable Pool + +We have introduced a variable pool in `Tool` to store variables, files, etc. generated during the tool's operation. These variables can be used by other tools during the tool's operation. + +Next, we will use `DallE3` and `Vectorizer.AI` as examples to introduce how to use the variable pool. + +- `DallE3` is an image generation tool that can generate images based on text. Here, we will let `DallE3` generate a logo for a coffee shop +- `Vectorizer.AI` is a vector image conversion tool that can convert images into vector images, so that the images can be infinitely enlarged without distortion. Here, we will convert the PNG icon generated by `DallE3` into a vector image, so that it can be truly used by designers. + +#### DallE3 +First, we use DallE3. After creating the image, we save the image to the variable pool. The code is as follows: + +```python +from typing import Any, Dict, List, Union +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +from base64 import b64decode + +from openai import OpenAI + +class DallE3Tool(BuiltinTool): + def _invoke(self, + user_id: str, + tool_parameters: Dict[str, Any], + ) -> Union[ToolInvokeMessage, List[ToolInvokeMessage]]: + """ + invoke tools + """ + client = OpenAI( + api_key=self.runtime.credentials['openai_api_key'], + ) + + # prompt + prompt = tool_parameters.get('prompt', '') + if not prompt: + return self.create_text_message('Please input prompt') + + # call openapi dalle3 + response = client.images.generate( + prompt=prompt, model='dall-e-3', + size='1024x1024', n=1, style='vivid', quality='standard', + response_format='b64_json' + ) + + result = [] + for image in response.data: + # Save all images to the variable pool through the save_as parameter. The variable name is self.VARIABLE_KEY.IMAGE.value. If new images are generated later, they will overwrite the previous images. + result.append(self.create_blob_message(blob=b64decode(image.b64_json), + meta={ 'mime_type': 'image/png' }, + save_as=self.VARIABLE_KEY.IMAGE.value)) + + return result +``` + +Note that we used `self.VARIABLE_KEY.IMAGE.value` as the variable name of the image. In order for developers' tools to cooperate with each other, we defined this `KEY`. You can use it freely, or you can choose not to use this `KEY`. Passing a custom KEY is also acceptable. + +#### Vectorizer.AI +Next, we use Vectorizer.AI to convert the PNG icon generated by DallE3 into a vector image. Let's go through the functions we defined here. The code is as follows: + +```python +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter +from core.tools.errors import ToolProviderCredentialValidationError + +from typing import Any, Dict, List, Union +from httpx import post +from base64 import b64decode + +class VectorizerTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: Dict[str, Any]) \ + -> Union[ToolInvokeMessage, List[ToolInvokeMessage]]: + """ + Tool invocation, the image variable name needs to be passed in from here, so that we can get the image from the variable pool + """ + + + def get_runtime_parameters(self) -> List[ToolParameter]: + """ + Override the tool parameter list, we can dynamically generate the parameter list based on the actual situation in the current variable pool, so that the LLM can generate the form based on the parameter list + """ + + + def is_tool_available(self) -> bool: + """ + Whether the current tool is available, if there is no image in the current variable pool, then we don't need to display this tool, just return False here + """ +``` + +Next, let's implement these three functions + +```python +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter +from core.tools.errors import ToolProviderCredentialValidationError + +from typing import Any, Dict, List, Union +from httpx import post +from base64 import b64decode + +class VectorizerTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: Dict[str, Any]) \ + -> Union[ToolInvokeMessage, List[ToolInvokeMessage]]: + """ + invoke tools + """ + api_key_name = self.runtime.credentials.get('api_key_name', None) + api_key_value = self.runtime.credentials.get('api_key_value', None) + + if not api_key_name or not api_key_value: + raise ToolProviderCredentialValidationError('Please input api key name and value') + + # Get image_id, the definition of image_id can be found in get_runtime_parameters + image_id = tool_parameters.get('image_id', '') + if not image_id: + return self.create_text_message('Please input image id') + + # Get the image generated by DallE from the variable pool + image_binary = self.get_variable_file(self.VARIABLE_KEY.IMAGE) + if not image_binary: + return self.create_text_message('Image not found, please request user to generate image firstly.') + + # Generate vector image + response = post( + 'https://vectorizer.ai/api/v1/vectorize', + files={ 'image': image_binary }, + data={ 'mode': 'test' }, + auth=(api_key_name, api_key_value), + timeout=30 + ) + + if response.status_code != 200: + raise Exception(response.text) + + return [ + self.create_text_message('the vectorized svg is saved as an image.'), + self.create_blob_message(blob=response.content, + meta={'mime_type': 'image/svg+xml'}) + ] + + def get_runtime_parameters(self) -> List[ToolParameter]: + """ + override the runtime parameters + """ + # Here, we override the tool parameter list, define the image_id, and set its option list to all images in the current variable pool. The configuration here is consistent with the configuration in yaml. + return [ + ToolParameter.get_simple_instance( + name='image_id', + llm_description=f'the image id that you want to vectorize, \ + and the image id should be specified in \ + {[i.name for i in self.list_default_image_variables()]}', + type=ToolParameter.ToolParameterType.SELECT, + required=True, + options=[i.name for i in self.list_default_image_variables()] + ) + ] + + def is_tool_available(self) -> bool: + # Only when there are images in the variable pool, the LLM needs to use this tool + return len(self.list_default_image_variables()) > 0 +``` + +It's worth noting that we didn't actually use `image_id` here. We assumed that there must be an image in the default variable pool when calling this tool, so we directly used `image_binary = self.get_variable_file(self.VARIABLE_KEY.IMAGE)` to get the image. In cases where the model's capabilities are weak, we recommend developers to do the same, which can effectively improve fault tolerance and avoid the model passing incorrect parameters. \ No newline at end of file diff --git a/api/core/tools/docs/en_US/tool_scale_out.md b/api/core/tools/docs/en_US/tool_scale_out.md new file mode 100644 index 0000000000000000000000000000000000000000..1deaf04a47539b68f466c4278583a934f80ea270 --- /dev/null +++ b/api/core/tools/docs/en_US/tool_scale_out.md @@ -0,0 +1,248 @@ +# Quick Tool Integration + +Here, we will use GoogleSearch as an example to demonstrate how to quickly integrate a tool. + +## 1. Prepare the Tool Provider yaml + +### Introduction + +This yaml declares a new tool provider, and includes information like the provider's name, icon, author, and other details that are fetched by the frontend for display. + +### Example + +We need to create a `google` module (folder) under `core/tools/provider/builtin`, and create `google.yaml`. The name must be consistent with the module name. + +Subsequently, all operations related to this tool will be carried out under this module. + +```yaml +identity: # Basic information of the tool provider + author: Dify # Author + name: google # Name, unique, no duplication with other providers + label: # Label for frontend display + en_US: Google # English label + zh_Hans: Google # Chinese label + description: # Description for frontend display + en_US: Google # English description + zh_Hans: Google # Chinese description + icon: icon.svg # Icon, needs to be placed in the _assets folder of the current module + tags: + - search + +``` + +- The `identity` field is mandatory, it contains the basic information of the tool provider, including author, name, label, description, icon, etc. + - The icon needs to be placed in the `_assets` folder of the current module, you can refer to [here](../../provider/builtin/google/_assets/icon.svg). + - The `tags` field is optional, it is used to classify the provider, and the frontend can filter the provider according to the tag, for all tags, they have been listed below: + + ```python + class ToolLabelEnum(Enum): + SEARCH = 'search' + IMAGE = 'image' + VIDEOS = 'videos' + WEATHER = 'weather' + FINANCE = 'finance' + DESIGN = 'design' + TRAVEL = 'travel' + SOCIAL = 'social' + NEWS = 'news' + MEDICAL = 'medical' + PRODUCTIVITY = 'productivity' + EDUCATION = 'education' + BUSINESS = 'business' + ENTERTAINMENT = 'entertainment' + UTILITIES = 'utilities' + OTHER = 'other' + ``` + +## 2. Prepare Provider Credentials + +Google, as a third-party tool, uses the API provided by SerpApi, which requires an API Key to use. This means that this tool needs a credential to use. For tools like `wikipedia`, there is no need to fill in the credential field, you can refer to [here](../../provider/builtin/wikipedia/wikipedia.yaml). + +After configuring the credential field, the effect is as follows: + +```yaml +identity: + author: Dify + name: google + label: + en_US: Google + zh_Hans: Google + description: + en_US: Google + zh_Hans: Google + icon: icon.svg +credentials_for_provider: # Credential field + serpapi_api_key: # Credential field name + type: secret-input # Credential field type + required: true # Required or not + label: # Credential field label + en_US: SerpApi API key # English label + zh_Hans: SerpApi API key # Chinese label + placeholder: # Credential field placeholder + en_US: Please input your SerpApi API key # English placeholder + zh_Hans: 请输入你的 SerpApi API key # Chinese placeholder + help: # Credential field help text + en_US: Get your SerpApi API key from SerpApi # English help text + zh_Hans: 从 SerpApi 获取您的 SerpApi API key # Chinese help text + url: https://serpapi.com/manage-api-key # Credential field help link + +``` + +- `type`: Credential field type, currently can be either `secret-input`, `text-input`, or `select` , corresponding to password input box, text input box, and drop-down box, respectively. If set to `secret-input`, it will mask the input content on the frontend, and the backend will encrypt the input content. + +## 3. Prepare Tool yaml + +A provider can have multiple tools, each tool needs a yaml file to describe, this file contains the basic information, parameters, output, etc. of the tool. + +Still taking GoogleSearch as an example, we need to create a `tools` module under the `google` module, and create `tools/google_search.yaml`, the content is as follows. + +```yaml +identity: # Basic information of the tool + name: google_search # Tool name, unique, no duplication with other tools + author: Dify # Author + label: # Label for frontend display + en_US: GoogleSearch # English label + zh_Hans: 谷歌搜索 # Chinese label +description: # Description for frontend display + human: # Introduction for frontend display, supports multiple languages + en_US: A tool for performing a Google SERP search and extracting snippets and webpages.Input should be a search query. + zh_Hans: 一个用于执行 Google SERP 搜索并提取片段和网页的工具。输入应该是一个搜索查询。 + llm: A tool for performing a Google SERP search and extracting snippets and webpages.Input should be a search query. # Introduction passed to LLM, in order to make LLM better understand this tool, we suggest to write as detailed information about this tool as possible here, so that LLM can understand and use this tool +parameters: # Parameter list + - name: query # Parameter name + type: string # Parameter type + required: true # Required or not + label: # Parameter label + en_US: Query string # English label + zh_Hans: 查询语句 # Chinese label + human_description: # Introduction for frontend display, supports multiple languages + en_US: used for searching + zh_Hans: 用于搜索网页内容 + llm_description: key words for searching # Introduction passed to LLM, similarly, in order to make LLM better understand this parameter, we suggest to write as detailed information about this parameter as possible here, so that LLM can understand this parameter + form: llm # Form type, llm means this parameter needs to be inferred by Agent, the frontend will not display this parameter + - name: result_type + type: select # Parameter type + required: true + options: # Drop-down box options + - value: text + label: + en_US: text + zh_Hans: 文本 + - value: link + label: + en_US: link + zh_Hans: 链接 + default: link + label: + en_US: Result type + zh_Hans: 结果类型 + human_description: + en_US: used for selecting the result type, text or link + zh_Hans: 用于选择结果类型,使用文本还是链接进行展示 + form: form # Form type, form means this parameter needs to be filled in by the user on the frontend before the conversation starts + +``` + +- The `identity` field is mandatory, it contains the basic information of the tool, including name, author, label, description, etc. +- `parameters` Parameter list + - `name` (Mandatory) Parameter name, must be unique and not duplicate with other parameters. + - `type` (Mandatory) Parameter type, currently supports `string`, `number`, `boolean`, `select`, `secret-input` five types, corresponding to string, number, boolean, drop-down box, and encrypted input box, respectively. For sensitive information, we recommend using the `secret-input` type + - `label` (Mandatory) Parameter label, for frontend display + - `form` (Mandatory) Form type, currently supports `llm`, `form` two types. + - In an agent app, `llm` indicates that the parameter is inferred by the LLM itself, while `form` indicates that the parameter can be pre-set for the tool. + - In a workflow app, both `llm` and `form` need to be filled out by the front end, but the parameters of `llm` will be used as input variables for the tool node. + - `required` Indicates whether the parameter is required or not + - In `llm` mode, if the parameter is required, the Agent is required to infer this parameter + - In `form` mode, if the parameter is required, the user is required to fill in this parameter on the frontend before the conversation starts + - `options` Parameter options + - In `llm` mode, Dify will pass all options to LLM, LLM can infer based on these options + - In `form` mode, when `type` is `select`, the frontend will display these options + - `default` Default value + - `min` Minimum value, can be set when the parameter type is `number`. + - `max` Maximum value, can be set when the parameter type is `number`. + - `placeholder` The prompt text for input boxes. It can be set when the form type is `form`, and the parameter type is `string`, `number`, or `secret-input`. It supports multiple languages. + - `human_description` Introduction for frontend display, supports multiple languages + - `llm_description` Introduction passed to LLM, in order to make LLM better understand this parameter, we suggest to write as detailed information about this parameter as possible here, so that LLM can understand this parameter + + +## 4. Add Tool Logic + +After completing the tool configuration, we can start writing the tool code that defines how it is invoked. + +Create `google_search.py` under the `google/tools` module, the content is as follows. + +```python +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.entities.tool_entities import ToolInvokeMessage + +from typing import Any, Dict, List, Union + +class GoogleSearchTool(BuiltinTool): + def _invoke(self, + user_id: str, + tool_parameters: Dict[str, Any], + ) -> Union[ToolInvokeMessage, List[ToolInvokeMessage]]: + """ + invoke tools + """ + query = tool_parameters['query'] + result_type = tool_parameters['result_type'] + api_key = self.runtime.credentials['serpapi_api_key'] + # Search with serpapi + result = SerpAPI(api_key).run(query, result_type=result_type) + + if result_type == 'text': + return self.create_text_message(text=result) + return self.create_link_message(link=result) +``` + +### Parameters + +The overall logic of the tool is in the `_invoke` method, this method accepts two parameters: `user_id` and `tool_parameters`, which represent the user ID and tool parameters respectively + +### Return Data + +When the tool returns, you can choose to return one message or multiple messages, here we return one message, using `create_text_message` and `create_link_message` can create a text message or a link message. If you want to return multiple messages, you can use `[self.create_text_message('msg1'), self.create_text_message('msg2')]` to create a list of messages. + +## 5. Add Provider Code + +Finally, we need to create a provider class under the provider module to implement the provider's credential verification logic. If the credential verification fails, it will throw a `ToolProviderCredentialValidationError` exception. + +Create `google.py` under the `google` module, the content is as follows. + +```python +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController +from core.tools.errors import ToolProviderCredentialValidationError + +from core.tools.provider.builtin.google.tools.google_search import GoogleSearchTool + +from typing import Any, Dict + +class GoogleProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: Dict[str, Any]) -> None: + try: + # 1. Here you need to instantiate a GoogleSearchTool with GoogleSearchTool(), it will automatically load the yaml configuration of GoogleSearchTool, but at this time it does not have credential information inside + # 2. Then you need to use the fork_tool_runtime method to pass the current credential information to GoogleSearchTool + # 3. Finally, invoke it, the parameters need to be passed according to the parameter rules configured in the yaml of GoogleSearchTool + GoogleSearchTool().fork_tool_runtime( + meta={ + "credentials": credentials, + } + ).invoke( + user_id='', + tool_parameters={ + "query": "test", + "result_type": "link" + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) +``` + +## Completion + +After the above steps are completed, we can see this tool on the frontend, and it can be used in the Agent. + +Of course, because google_search needs a credential, before using it, you also need to input your credentials on the frontend. + +![Alt text](../images/index/image-2.png) diff --git a/api/core/tools/docs/images/index/image-1.png b/api/core/tools/docs/images/index/image-1.png new file mode 100644 index 0000000000000000000000000000000000000000..20da5569faa2282eb9d93cf306a1142456ce4caa --- /dev/null +++ b/api/core/tools/docs/images/index/image-1.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0103ec75e6033f5f24434f0ff2d594be400b7617418201533c39d7ed1e31b0b3 +size 248210 diff --git a/api/core/tools/docs/images/index/image-2.png b/api/core/tools/docs/images/index/image-2.png new file mode 100644 index 0000000000000000000000000000000000000000..eaab65938b1858e163228453fb281c100e240793 --- /dev/null +++ b/api/core/tools/docs/images/index/image-2.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6939e5b06191bcb57df0e44661bcbe35f61acc73e053f748c021265881233488 +size 416650 diff --git a/api/core/tools/docs/images/index/image.png b/api/core/tools/docs/images/index/image.png new file mode 100644 index 0000000000000000000000000000000000000000..ef01009ff089ffb433b208c3796e4fd5d1f55a6a --- /dev/null +++ b/api/core/tools/docs/images/index/image.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bc6c88e23532b6298ccdf63724caf40c950a7434bc9a8ab91e9e0b4600ae375e +size 272811 diff --git a/api/core/tools/docs/ja_JP/advanced_scale_out.md b/api/core/tools/docs/ja_JP/advanced_scale_out.md new file mode 100644 index 0000000000000000000000000000000000000000..96f843354f91b535c6eefdc8612a62b1d00782c1 --- /dev/null +++ b/api/core/tools/docs/ja_JP/advanced_scale_out.md @@ -0,0 +1,283 @@ +# 高度なツール統合 + +このガイドを始める前に、Difyのツール統合プロセスの基本を理解していることを確認してください。簡単な概要については[クイック統合](./tool_scale_out.md)をご覧ください。 + +## ツールインターフェース + +より複雑なツールを迅速に構築するのを支援するため、`Tool`クラスに一連のヘルパーメソッドを定義しています。 + +### メッセージの返却 + +Difyは`テキスト`、`リンク`、`画像`、`ファイルBLOB`、`JSON`などの様々なメッセージタイプをサポートしています。以下のインターフェースを通じて、異なるタイプのメッセージをLLMとユーザーに返すことができます。 + +注意:以下のインターフェースの一部のパラメータについては、後のセクションで説明します。 + +#### 画像URL +画像のURLを渡すだけで、Difyが自動的に画像をダウンロードしてユーザーに返します。 + +```python + def create_image_message(self, image: str, save_as: str = '') -> ToolInvokeMessage: + """ + create an image message + + :param image: the url of the image + :param save_as: save as + :return: the image message + """ +``` + +#### リンク +リンクを返す必要がある場合は、以下のインターフェースを使用できます。 + +```python + def create_link_message(self, link: str, save_as: str = '') -> ToolInvokeMessage: + """ + create a link message + + :param link: the url of the link + :param save_as: save as + :return: the link message + """ +``` + +#### テキスト +テキストメッセージを返す必要がある場合は、以下のインターフェースを使用できます。 + +```python + def create_text_message(self, text: str, save_as: str = '') -> ToolInvokeMessage: + """ + create a text message + + :param text: the text of the message + :param save_as: save as + :return: the text message + """ +``` + +#### ファイルBLOB +画像、音声、動画、PPT、Word、Excelなどのファイルの生データを返す必要がある場合は、以下のインターフェースを使用できます。 + +- `blob` ファイルの生データ(bytes型) +- `meta` ファイルのメタデータ。ファイルの種類が分かっている場合は、`mime_type`を渡すことをお勧めします。そうでない場合、Difyはデフォルトタイプとして`octet/stream`を使用します。 + +```python + def create_blob_message(self, blob: bytes, meta: dict = None, save_as: str = '') -> ToolInvokeMessage: + """ + create a blob message + + :param blob: the blob + :param meta: meta + :param save_as: save as + :return: the blob message + """ +``` + +#### JSON +フォーマットされたJSONを返す必要がある場合は、以下のインターフェースを使用できます。これは通常、ワークフロー内のノード間のデータ伝送に使用されますが、エージェントモードでは、ほとんどの大規模言語モデルもJSONを読み取り、理解することができます。 + +- `object` Pythonの辞書オブジェクトで、自動的にJSONにシリアライズされます。 + +```python + def create_json_message(self, object: dict) -> ToolInvokeMessage: + """ + create a json message + """ +``` + +### ショートカットツール + +大規模モデルアプリケーションでは、以下の2つの一般的なニーズがあります: +- まず長いテキストを事前に要約し、その要約内容をLLMに渡すことで、元のテキストが長すぎてLLMが処理できない問題を防ぐ +- ツールが取得したコンテンツがリンクである場合、Webページ情報をクロールしてからLLMに返す必要がある + +開発者がこれら2つのニーズを迅速に実装できるよう、以下の2つのショートカットツールを提供しています。 + +#### テキスト要約ツール + +このツールはuser_idと要約するテキストを入力として受け取り、要約されたテキストを返します。Difyは現在のワークスペースのデフォルトモデルを使用して長文を要約します。 + +```python + def summary(self, user_id: str, content: str) -> str: + """ + summary the content + + :param user_id: the user id + :param content: the content + :return: the summary + """ +``` + +#### Webページクローリングツール + +このツールはクロールするWebページのリンクとユーザーエージェント(空でも可)を入力として受け取り、そのWebページの情報を含む文字列を返します。`user_agent`はオプションのパラメータで、ツールを識別するために使用できます。渡さない場合、Difyはデフォルトの`user_agent`を使用します。 + +```python + def get_url(self, url: str, user_agent: str = None) -> str: + """ + get url from the crawled result + """ +``` + +### 変数プール + +`Tool`内に変数プールを導入し、ツールの実行中に生成された変数やファイルなどを保存します。これらの変数は、ツールの実行中に他のツールが使用することができます。 + +次に、`DallE3`と`Vectorizer.AI`を例に、変数プールの使用方法を紹介します。 + +- `DallE3`は画像生成ツールで、テキストに基づいて画像を生成できます。ここでは、`DallE3`にカフェのロゴを生成させます。 +- `Vectorizer.AI`はベクター画像変換ツールで、画像をベクター画像に変換できるため、画像を無限に拡大しても品質が損なわれません。ここでは、`DallE3`が生成したPNGアイコンをベクター画像に変換し、デザイナーが実際に使用できるようにします。 + +#### DallE3 +まず、DallE3を使用します。画像を作成した後、その画像を変数プールに保存します。コードは以下の通りです: + +```python +from typing import Any, Dict, List, Union +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +from base64 import b64decode + +from openai import OpenAI + +class DallE3Tool(BuiltinTool): + def _invoke(self, + user_id: str, + tool_parameters: Dict[str, Any], + ) -> Union[ToolInvokeMessage, List[ToolInvokeMessage]]: + """ + invoke tools + """ + client = OpenAI( + api_key=self.runtime.credentials['openai_api_key'], + ) + + # prompt + prompt = tool_parameters.get('prompt', '') + if not prompt: + return self.create_text_message('Please input prompt') + + # call openapi dalle3 + response = client.images.generate( + prompt=prompt, model='dall-e-3', + size='1024x1024', n=1, style='vivid', quality='standard', + response_format='b64_json' + ) + + result = [] + for image in response.data: + # Save all images to the variable pool through the save_as parameter. The variable name is self.VARIABLE_KEY.IMAGE.value. If new images are generated later, they will overwrite the previous images. + result.append(self.create_blob_message(blob=b64decode(image.b64_json), + meta={ 'mime_type': 'image/png' }, + save_as=self.VARIABLE_KEY.IMAGE.value)) + + return result +``` + +ここでは画像の変数名として`self.VARIABLE_KEY.IMAGE.value`を使用していることに注意してください。開発者のツールが互いに連携できるよう、この`KEY`を定義しました。自由に使用することも、この`KEY`を使用しないこともできます。カスタムのKEYを渡すこともできます。 + +#### Vectorizer.AI +次に、Vectorizer.AIを使用して、DallE3が生成したPNGアイコンをベクター画像に変換します。ここで定義した関数を見てみましょう。コードは以下の通りです: + +```python +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter +from core.tools.errors import ToolProviderCredentialValidationError + +from typing import Any, Dict, List, Union +from httpx import post +from base64 import b64decode + +class VectorizerTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: Dict[str, Any]) + -> Union[ToolInvokeMessage, List[ToolInvokeMessage]]: + """ + Tool invocation, the image variable name needs to be passed in from here, so that we can get the image from the variable pool + """ + + + def get_runtime_parameters(self) -> List[ToolParameter]: + """ + Override the tool parameter list, we can dynamically generate the parameter list based on the actual situation in the current variable pool, so that the LLM can generate the form based on the parameter list + """ + + + def is_tool_available(self) -> bool: + """ + Whether the current tool is available, if there is no image in the current variable pool, then we don't need to display this tool, just return False here + """ +``` + +次に、これら3つの関数を実装します: + +```python +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter +from core.tools.errors import ToolProviderCredentialValidationError + +from typing import Any, Dict, List, Union +from httpx import post +from base64 import b64decode + +class VectorizerTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: Dict[str, Any]) + -> Union[ToolInvokeMessage, List[ToolInvokeMessage]]: + """ + invoke tools + """ + api_key_name = self.runtime.credentials.get('api_key_name', None) + api_key_value = self.runtime.credentials.get('api_key_value', None) + + if not api_key_name or not api_key_value: + raise ToolProviderCredentialValidationError('Please input api key name and value') + + # Get image_id, the definition of image_id can be found in get_runtime_parameters + image_id = tool_parameters.get('image_id', '') + if not image_id: + return self.create_text_message('Please input image id') + + # Get the image generated by DallE from the variable pool + image_binary = self.get_variable_file(self.VARIABLE_KEY.IMAGE) + if not image_binary: + return self.create_text_message('Image not found, please request user to generate image firstly.') + + # Generate vector image + response = post( + 'https://vectorizer.ai/api/v1/vectorize', + files={ 'image': image_binary }, + data={ 'mode': 'test' }, + auth=(api_key_name, api_key_value), + timeout=30 + ) + + if response.status_code != 200: + raise Exception(response.text) + + return [ + self.create_text_message('the vectorized svg is saved as an image.'), + self.create_blob_message(blob=response.content, + meta={'mime_type': 'image/svg+xml'}) + ] + + def get_runtime_parameters(self) -> List[ToolParameter]: + """ + override the runtime parameters + """ + # Here, we override the tool parameter list, define the image_id, and set its option list to all images in the current variable pool. The configuration here is consistent with the configuration in yaml. + return [ + ToolParameter.get_simple_instance( + name='image_id', + llm_description=f'the image id that you want to vectorize, \ + and the image id should be specified in \ + {[i.name for i in self.list_default_image_variables()]}', + type=ToolParameter.ToolParameterType.SELECT, + required=True, + options=[i.name for i in self.list_default_image_variables()] + ) + ] + + def is_tool_available(self) -> bool: + # Only when there are images in the variable pool, the LLM needs to use this tool + return len(self.list_default_image_variables()) > 0 +``` + +ここで注目すべきは、実際には`image_id`を使用していないことです。このツールを呼び出す際には、デフォルトの変数プールに必ず画像があると仮定し、直接`image_binary = self.get_variable_file(self.VARIABLE_KEY.IMAGE)`を使用して画像を取得しています。モデルの能力が弱い場合、開発者にもこの方法を推奨します。これにより、エラー許容度を効果的に向上させ、モデルが誤ったパラメータを渡すのを防ぐことができます。 \ No newline at end of file diff --git a/api/core/tools/docs/ja_JP/tool_scale_out.md b/api/core/tools/docs/ja_JP/tool_scale_out.md new file mode 100644 index 0000000000000000000000000000000000000000..a721023d00bdda49569896a84494752ecf16a2f1 --- /dev/null +++ b/api/core/tools/docs/ja_JP/tool_scale_out.md @@ -0,0 +1,240 @@ +# ツールの迅速な統合 + +ここでは、GoogleSearchを例にツールを迅速に統合する方法を紹介します。 + +## 1. ツールプロバイダーのyamlを準備する + +### 概要 + +このyamlファイルには、プロバイダー名、アイコン、作者などの詳細情報が含まれ、フロントエンドでの柔軟な表示を可能にします。 + +### 例 + +`core/tools/provider/builtin`の下に`google`モジュール(フォルダ)を作成し、`google.yaml`を作成します。名前はモジュール名と一致している必要があります。 + +以降、このツールに関するすべての操作はこのモジュール内で行います。 + +```yaml +identity: # ツールプロバイダーの基本情報 + author: Dify # 作者 + name: google # 名前(一意、他のプロバイダーと重複不可) + label: # フロントエンド表示用のラベル + en_US: Google # 英語ラベル + zh_Hans: Google # 中国語ラベル + description: # フロントエンド表示用の説明 + en_US: Google # 英語説明 + zh_Hans: Google # 中国語説明 + icon: icon.svg # アイコン(現在のモジュールの_assetsフォルダに配置) + tags: # タグ(フロントエンド表示用) + - search +``` + +- `identity`フィールドは必須で、ツールプロバイダーの基本情報(作者、名前、ラベル、説明、アイコンなど)が含まれます。 + - アイコンは現在のモジュールの`_assets`フォルダに配置する必要があります。[こちら](../../provider/builtin/google/_assets/icon.svg)を参照してください。 + - タグはフロントエンドでの表示に使用され、ユーザーがこのツールプロバイダーを素早く見つけるのに役立ちます。現在サポートされているすべてのタグは以下の通りです: + ```python + class ToolLabelEnum(Enum): + SEARCH = 'search' + IMAGE = 'image' + VIDEOS = 'videos' + WEATHER = 'weather' + FINANCE = 'finance' + DESIGN = 'design' + TRAVEL = 'travel' + SOCIAL = 'social' + NEWS = 'news' + MEDICAL = 'medical' + PRODUCTIVITY = 'productivity' + EDUCATION = 'education' + BUSINESS = 'business' + ENTERTAINMENT = 'entertainment' + UTILITIES = 'utilities' + OTHER = 'other' + ``` + +## 2. プロバイダーの認証情報を準備する + +GoogleはSerpApiが提供するAPIを使用するサードパーティツールであり、SerpApiを使用するにはAPI Keyが必要です。つまり、このツールを使用するには認証情報が必要です。一方、`wikipedia`のようなツールでは認証情報フィールドを記入する必要はありません。[こちら](../../provider/builtin/wikipedia/wikipedia.yaml)を参照してください。 + +認証情報フィールドを設定すると、以下のようになります: + +```yaml +identity: + author: Dify + name: google + label: + en_US: Google + zh_Hans: Google + description: + en_US: Google + zh_Hans: Google + icon: icon.svg +credentials_for_provider: # 認証情報フィールド + serpapi_api_key: # 認証情報フィールド名 + type: secret-input # 認証情報フィールドタイプ + required: true # 必須かどうか + label: # 認証情報フィールドラベル + en_US: SerpApi API key # 英語ラベル + zh_Hans: SerpApi API key # 中国語ラベル + placeholder: # 認証情報フィールドプレースホルダー + en_US: Please input your SerpApi API key # 英語プレースホルダー + zh_Hans: 请输入你的 SerpApi API key # 中国語プレースホルダー + help: # 認証情報フィールドヘルプテキスト + en_US: Get your SerpApi API key from SerpApi # 英語ヘルプテキスト + zh_Hans: 从 SerpApi 获取您的 SerpApi API key # 中国語ヘルプテキスト + url: https://serpapi.com/manage-api-key # 認証情報フィールドヘルプリンク +``` + +- `type`:認証情報フィールドタイプ。現在、`secret-input`、`text-input`、`select`の3種類をサポートしており、それぞれパスワード入力ボックス、テキスト入力ボックス、ドロップダウンボックスに対応します。`secret-input`の場合、フロントエンドで入力内容が隠され、バックエンドで入力内容が暗号化されます。 + +## 3. ツールのyamlを準備する + +1つのプロバイダーの下に複数のツールを持つことができ、各ツールにはyamlファイルが必要です。このファイルにはツールの基本情報、パラメータ、出力などが含まれます。 + +引き続きGoogleSearchを例に、`google`モジュールの下に`tools`モジュールを作成し、`tools/google_search.yaml`を作成します。内容は以下の通りです: + +```yaml +identity: # ツールの基本情報 + name: google_search # ツール名(一意、他のツールと重複不可) + author: Dify # 作者 + label: # フロントエンド表示用のラベル + en_US: GoogleSearch # 英語ラベル + zh_Hans: 谷歌搜索 # 中国語ラベル +description: # フロントエンド表示用の説明 + human: # フロントエンド表示用の紹介(多言語対応) + en_US: A tool for performing a Google SERP search and extracting snippets and webpages. Input should be a search query. + zh_Hans: 一个用于执行 Google SERP 搜索并提取片段和网页的工具。输入应该是一个搜索查询。 + llm: A tool for performing a Google SERP search and extracting snippets and webpages. Input should be a search query. # LLMに渡す紹介文。LLMがこのツールをより理解できるよう、できるだけ詳細な情報を記述することをお勧めします。 +parameters: # パラメータリスト + - name: query # パラメータ名 + type: string # パラメータタイプ + required: true # 必須かどうか + label: # パラメータラベル + en_US: Query string # 英語ラベル + zh_Hans: 查询语句 # 中国語ラベル + human_description: # フロントエンド表示用の紹介(多言語対応) + en_US: used for searching + zh_Hans: 用于搜索网页内容 + llm_description: key words for searching # LLMに渡す紹介文。LLMがこのパラメータをより理解できるよう、できるだけ詳細な情報を記述することをお勧めします。 + form: llm # フォームタイプ。llmはこのパラメータがAgentによって推論される必要があることを意味し、フロントエンドではこのパラメータは表示されません。 + - name: result_type + type: select # パラメータタイプ + required: true + options: # ドロップダウンボックスのオプション + - value: text + label: + en_US: text + zh_Hans: 文本 + - value: link + label: + en_US: link + zh_Hans: 链接 + default: link + label: + en_US: Result type + zh_Hans: 结果类型 + human_description: + en_US: used for selecting the result type, text or link + zh_Hans: 用于选择结果类型,使用文本还是链接进行展示 + form: form # フォームタイプ。formはこのパラメータが対話開始前にフロントエンドでユーザーによって入力される必要があることを意味します。 +``` + +- `identity`フィールドは必須で、ツールの基本情報(名前、作者、ラベル、説明など)が含まれます。 +- `parameters` パラメータリスト + - `name`(必須)パラメータ名。一意で、他のパラメータと重複しないようにしてください。 + - `type`(必須)パラメータタイプ。現在、`string`、`number`、`boolean`、`select`、`secret-input`の5種類をサポートしており、それぞれ文字列、数値、ブール値、ドロップダウンボックス、暗号化入力ボックスに対応します。機密情報には`secret-input`タイプの使用をお勧めします。 + - `label`(必須)パラメータラベル。フロントエンド表示用です。 + - `form`(必須)フォームタイプ。現在、`llm`と`form`の2種類をサポートしています。 + - エージェントアプリケーションでは、`llm`はこのパラメータがLLM自身によって推論されることを示し、`form`はこのツールを使用するために事前に設定できるパラメータであることを示します。 + - ワークフローアプリケーションでは、`llm`と`form`の両方がフロントエンドで入力する必要がありますが、`llm`のパラメータはツールノードの入力変数として使用されます。 + - `required` パラメータが必須かどうかを示します。 + - `llm`モードでは、パラメータが必須の場合、Agentはこのパラメータを推論する必要があります。 + - `form`モードでは、パラメータが必須の場合、ユーザーは対話開始前にフロントエンドでこのパラメータを入力する必要があります。 + - `options` パラメータオプション + - `llm`モードでは、DifyはすべてのオプションをLLMに渡し、LLMはこれらのオプションに基づいて推論できます。 + - `form`モードで、`type`が`select`の場合、フロントエンドはこれらのオプションを表示します。 + - `default` デフォルト値 + - `min` 最小値。パラメータタイプが`number`の場合に設定できます。 + - `max` 最大値。パラメータタイプが`number`の場合に設定できます。 + - `human_description` フロントエンド表示用の紹介。多言語対応です。 + - `placeholder` 入力ボックスのプロンプトテキスト。フォームタイプが`form`で、パラメータタイプが`string`、`number`、`secret-input`の場合に設定できます。多言語対応です。 + - `llm_description` LLMに渡す紹介文。LLMがこのパラメータをより理解できるよう、できるだけ詳細な情報を記述することをお勧めします。 + +## 4. ツールコードを準備する + +ツールの設定が完了したら、ツールのロジックを実装するコードを作成します。 + +`google/tools`モジュールの下に`google_search.py`を作成し、内容は以下の通りです: + +```python +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.entities.tool_entities import ToolInvokeMessage + +from typing import Any, Dict, List, Union + +class GoogleSearchTool(BuiltinTool): + def _invoke(self, + user_id: str, + tool_parameters: Dict[str, Any], + ) -> Union[ToolInvokeMessage, List[ToolInvokeMessage]]: + """ + ツールを呼び出す + """ + query = tool_parameters['query'] + result_type = tool_parameters['result_type'] + api_key = self.runtime.credentials['serpapi_api_key'] + result = SerpAPI(api_key).run(query, result_type=result_type) + + if result_type == 'text': + return self.create_text_message(text=result) + return self.create_link_message(link=result) +``` + +### パラメータ +ツールの全体的なロジックは`_invoke`メソッドにあります。このメソッドは2つのパラメータ(`user_id`とtool_parameters`)を受け取り、それぞれユーザーIDとツールパラメータを表します。 + +### 戻り値 +ツールの戻り値として、1つのメッセージまたは複数のメッセージを選択できます。ここでは1つのメッセージを返しています。`create_text_message`と`create_link_message`を使用して、テキストメッセージまたはリンクメッセージを作成できます。複数のメッセージを返す場合は、リストを構築できます(例:`[self.create_text_message('msg1'), self.create_text_message('msg2')]`)。 + +## 5. プロバイダーコードを準備する + +最後に、プロバイダーモジュールの下にプロバイダークラスを作成し、プロバイダーの認証情報検証ロジックを実装する必要があります。認証情報の検証が失敗した場合、`ToolProviderCredentialValidationError`例外が発生します。 + +`google`モジュールの下に`google.py`を作成し、内容は以下の通りです: + +```python +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController +from core.tools.errors import ToolProviderCredentialValidationError + +from core.tools.provider.builtin.google.tools.google_search import GoogleSearchTool + +from typing import Any, Dict + +class GoogleProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: Dict[str, Any]) -> None: + try: + # 1. ここでGoogleSearchTool()を使ってGoogleSearchToolをインスタンス化する必要があります。これによりGoogleSearchToolのyaml設定が自動的に読み込まれますが、この時点では認証情報は含まれていません + # 2. 次に、fork_tool_runtimeメソッドを使用して、現在の認証情報をGoogleSearchToolに渡す必要があります + # 3. 最後に、invokeを呼び出します。パラメータはGoogleSearchToolのyamlで設定されたパラメータルールに従って渡す必要があります + GoogleSearchTool().fork_tool_runtime( + meta={ + "credentials": credentials, + } + ).invoke( + user_id='', + tool_parameters={ + "query": "test", + "result_type": "link" + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) +``` + +## 完了 + +以上のステップが完了すると、このツールをフロントエンドで確認し、Agentで使用することができるようになります。 + +もちろん、google_searchには認証情報が必要なため、使用する前にフロントエンドで認証情報を入力する必要があります。 + +![Alt text](../images/index/image-2.png) \ No newline at end of file diff --git a/api/core/tools/docs/zh_Hans/advanced_scale_out.md b/api/core/tools/docs/zh_Hans/advanced_scale_out.md new file mode 100644 index 0000000000000000000000000000000000000000..0385dfe4e7bce0598debbd679c63b3ede8c6dfa3 --- /dev/null +++ b/api/core/tools/docs/zh_Hans/advanced_scale_out.md @@ -0,0 +1,283 @@ +# 高级接入Tool + +在开始高级接入之前,请确保你已经阅读过[快速接入](./tool_scale_out.md),并对Dify的工具接入流程有了基本的了解。 + +## 工具接口 + +我们在`Tool`类中定义了一系列快捷方法,用于帮助开发者快速构较为复杂的工具 + +### 消息返回 + +Dify支持`文本` `链接` `图片` `文件BLOB` `JSON` 等多种消息类型,你可以通过以下几个接口返回不同类型的消息给LLM和用户。 + +注意,在下面的接口中的部分参数将在后面的章节中介绍。 + +#### 图片URL +只需要传递图片的URL即可,Dify会自动下载图片并返回给用户。 + +```python + def create_image_message(self, image: str, save_as: str = '') -> ToolInvokeMessage: + """ + create an image message + + :param image: the url of the image + :param save_as: save as + :return: the image message + """ +``` + +#### 链接 +如果你需要返回一个链接,可以使用以下接口。 + +```python + def create_link_message(self, link: str, save_as: str = '') -> ToolInvokeMessage: + """ + create a link message + + :param link: the url of the link + :param save_as: save as + :return: the link message + """ +``` + +#### 文本 +如果你需要返回一个文本消息,可以使用以下接口。 + +```python + def create_text_message(self, text: str, save_as: str = '') -> ToolInvokeMessage: + """ + create a text message + + :param text: the text of the message + :param save_as: save as + :return: the text message + """ +``` + +#### 文件BLOB +如果你需要返回文件的原始数据,如图片、音频、视频、PPT、Word、Excel等,可以使用以下接口。 + +- `blob` 文件的原始数据,bytes类型 +- `meta` 文件的元数据,如果你知道该文件的类型,最好传递一个`mime_type`,否则Dify将使用`octet/stream`作为默认类型 + +```python + def create_blob_message(self, blob: bytes, meta: dict = None, save_as: str = '') -> ToolInvokeMessage: + """ + create a blob message + + :param blob: the blob + :param meta: meta + :param save_as: save as + :return: the blob message + """ +``` + +#### JSON +如果你需要返回一个格式化的JSON,可以使用以下接口。这通常用于workflow中的节点间的数据传递,当然agent模式中,大部分大模型也都能够阅读和理解JSON。 + +- `object` 一个Python的字典对象,会被自动序列化为JSON + +```python + def create_json_message(self, object: dict) -> ToolInvokeMessage: + """ + create a json message + """ +``` + +### 快捷工具 + +在大模型应用中,我们有两种常见的需求: +- 先将很长的文本进行提前总结,然后再将总结内容传递给LLM,以防止原文本过长导致LLM无法处理 +- 工具获取到的内容是一个链接,需要爬取网页信息后再返回给LLM + +为了帮助开发者快速实现这两种需求,我们提供了以下两个快捷工具。 + +#### 文本总结工具 + +该工具需要传入user_id和需要进行总结的文本,返回一个总结后的文本,Dify会使用当前工作空间的默认模型对长文本进行总结。 + +```python + def summary(self, user_id: str, content: str) -> str: + """ + summary the content + + :param user_id: the user id + :param content: the content + :return: the summary + """ +``` + +#### 网页爬取工具 + +该工具需要传入需要爬取的网页链接和一个user_agent(可为空),返回一个包含该网页信息的字符串,其中`user_agent`是可选参数,可以用来识别工具,如果不传递,Dify将使用默认的`user_agent`。 + +```python + def get_url(self, url: str, user_agent: str = None) -> str: + """ + get url from the crawled result + """ +``` + +### 变量池 + +我们在`Tool`中引入了一个变量池,用于存储工具运行过程中产生的变量、文件等,这些变量可以在工具运行过程中被其他工具使用。 + +下面,我们以`DallE3`和`Vectorizer.AI`为例,介绍如何使用变量池。 + +- `DallE3`是一个图片生成工具,它可以根据文本生成图片,在这里,我们将让`DallE3`生成一个咖啡厅的Logo +- `Vectorizer.AI`是一个矢量图转换工具,它可以将图片转换为矢量图,使得图片可以无限放大而不失真,在这里,我们将`DallE3`生成的PNG图标转换为矢量图,从而可以真正被设计师使用。 + +#### DallE3 +首先我们使用DallE3,在创建完图片以后,我们将图片保存到变量池中,代码如下 + +```python +from typing import Any, Dict, List, Union +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +from base64 import b64decode + +from openai import OpenAI + +class DallE3Tool(BuiltinTool): + def _invoke(self, + user_id: str, + tool_parameters: Dict[str, Any], + ) -> Union[ToolInvokeMessage, List[ToolInvokeMessage]]: + """ + invoke tools + """ + client = OpenAI( + api_key=self.runtime.credentials['openai_api_key'], + ) + + # prompt + prompt = tool_parameters.get('prompt', '') + if not prompt: + return self.create_text_message('Please input prompt') + + # call openapi dalle3 + response = client.images.generate( + prompt=prompt, model='dall-e-3', + size='1024x1024', n=1, style='vivid', quality='standard', + response_format='b64_json' + ) + + result = [] + for image in response.data: + # 将所有图片通过save_as参数保存到变量池中,变量名为self.VARIABLE_KEY.IMAGE.value,如果如果后续有新的图片生成,那么将会覆盖之前的图片 + result.append(self.create_blob_message(blob=b64decode(image.b64_json), + meta={ 'mime_type': 'image/png' }, + save_as=self.VARIABLE_KEY.IMAGE.value)) + + return result +``` + +我们可以注意到这里我们使用了`self.VARIABLE_KEY.IMAGE.value`作为图片的变量名,为了便于开发者们的工具能够互相配合,我们定义了这个`KEY`,大家可以自由使用,也可以不使用这个`KEY`,传递一个自定义的KEY也是可以的。 + +#### Vectorizer.AI +接下来我们使用Vectorizer.AI,将DallE3生成的PNG图标转换为矢量图,我们先来过一遍我们在这里定义的函数,代码如下 + +```python +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter +from core.tools.errors import ToolProviderCredentialValidationError + +from typing import Any, Dict, List, Union +from httpx import post +from base64 import b64decode + +class VectorizerTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: Dict[str, Any]) \ + -> Union[ToolInvokeMessage, List[ToolInvokeMessage]]: + """ + 工具调用,图片变量名需要从这里传递进来,从而我们就可以从变量池中获取到图片 + """ + + + def get_runtime_parameters(self) -> List[ToolParameter]: + """ + 重写工具参数列表,我们可以根据当前变量池里的实际情况来动态生成参数列表,从而LLM可以根据参数列表来生成表单 + """ + + + def is_tool_available(self) -> bool: + """ + 当前工具是否可用,如果当前变量池中没有图片,那么我们就不需要展示这个工具,这里返回False即可 + """ +``` + +接下来我们来实现这三个函数 + +```python +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter +from core.tools.errors import ToolProviderCredentialValidationError + +from typing import Any, Dict, List, Union +from httpx import post +from base64 import b64decode + +class VectorizerTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: Dict[str, Any]) \ + -> Union[ToolInvokeMessage, List[ToolInvokeMessage]]: + """ + invoke tools + """ + api_key_name = self.runtime.credentials.get('api_key_name', None) + api_key_value = self.runtime.credentials.get('api_key_value', None) + + if not api_key_name or not api_key_value: + raise ToolProviderCredentialValidationError('Please input api key name and value') + + # 获取image_id,image_id的定义可以在get_runtime_parameters中找到 + image_id = tool_parameters.get('image_id', '') + if not image_id: + return self.create_text_message('Please input image id') + + # 从变量池中获取到之前DallE生成的图片 + image_binary = self.get_variable_file(self.VARIABLE_KEY.IMAGE) + if not image_binary: + return self.create_text_message('Image not found, please request user to generate image firstly.') + + # 生成矢量图 + response = post( + 'https://vectorizer.ai/api/v1/vectorize', + files={ 'image': image_binary }, + data={ 'mode': 'test' }, + auth=(api_key_name, api_key_value), + timeout=30 + ) + + if response.status_code != 200: + raise Exception(response.text) + + return [ + self.create_text_message('the vectorized svg is saved as an image.'), + self.create_blob_message(blob=response.content, + meta={'mime_type': 'image/svg+xml'}) + ] + + def get_runtime_parameters(self) -> List[ToolParameter]: + """ + override the runtime parameters + """ + # 这里,我们重写了工具参数列表,定义了image_id,并设置了它的选项列表为当前变量池中的所有图片,这里的配置与yaml中的配置是一致的 + return [ + ToolParameter.get_simple_instance( + name='image_id', + llm_description=f'the image id that you want to vectorize, \ + and the image id should be specified in \ + {[i.name for i in self.list_default_image_variables()]}', + type=ToolParameter.ToolParameterType.SELECT, + required=True, + options=[i.name for i in self.list_default_image_variables()] + ) + ] + + def is_tool_available(self) -> bool: + # 只有当变量池中有图片时,LLM才需要使用这个工具 + return len(self.list_default_image_variables()) > 0 +``` + +可以注意到的是,我们这里其实并没有使用到`image_id`,我们已经假设了调用这个工具的时候一定有一张图片在默认的变量池中,所以直接使用了`image_binary = self.get_variable_file(self.VARIABLE_KEY.IMAGE)`来获取图片,在模型能力较弱的情况下,我们建议开发者们也这样做,可以有效提升容错率,避免模型传递错误的参数。 \ No newline at end of file diff --git a/api/core/tools/docs/zh_Hans/tool_scale_out.md b/api/core/tools/docs/zh_Hans/tool_scale_out.md new file mode 100644 index 0000000000000000000000000000000000000000..ec61e4677bae76d3c11cf948fa860732eae811bf --- /dev/null +++ b/api/core/tools/docs/zh_Hans/tool_scale_out.md @@ -0,0 +1,237 @@ +# 快速接入Tool + +这里我们以GoogleSearch为例,介绍如何快速接入一个工具。 + +## 1. 准备工具供应商yaml + +### 介绍 +这个yaml将包含工具供应商的信息,包括供应商名称、图标、作者等详细信息,以帮助前端灵活展示。 + +### 示例 + +我们需要在 `core/tools/provider/builtin`下创建一个`google`模块(文件夹),并创建`google.yaml`,名称必须与模块名称一致。 + +后续,我们关于这个工具的所有操作都将在这个模块下进行。 + +```yaml +identity: # 工具供应商的基本信息 + author: Dify # 作者 + name: google # 名称,唯一,不允许和其他供应商重名 + label: # 标签,用于前端展示 + en_US: Google # 英文标签 + zh_Hans: Google # 中文标签 + description: # 描述,用于前端展示 + en_US: Google # 英文描述 + zh_Hans: Google # 中文描述 + icon: icon.svg # 图标,需要放置在当前模块的_assets文件夹下 + tags: # 标签,用于前端展示 + - search + +``` + - `identity` 字段是必须的,它包含了工具供应商的基本信息,包括作者、名称、标签、描述、图标等 + - 图标需要放置在当前模块的`_assets`文件夹下,可以参考[这里](../../provider/builtin/google/_assets/icon.svg)。 + - 标签用于前端展示,可以帮助用户快速找到这个工具供应商,下面列出了目前所支持的所有标签 + ```python + class ToolLabelEnum(Enum): + SEARCH = 'search' + IMAGE = 'image' + VIDEOS = 'videos' + WEATHER = 'weather' + FINANCE = 'finance' + DESIGN = 'design' + TRAVEL = 'travel' + SOCIAL = 'social' + NEWS = 'news' + MEDICAL = 'medical' + PRODUCTIVITY = 'productivity' + EDUCATION = 'education' + BUSINESS = 'business' + ENTERTAINMENT = 'entertainment' + UTILITIES = 'utilities' + OTHER = 'other' + ``` + +## 2. 准备供应商凭据 + +Google作为一个第三方工具,使用了SerpApi提供的API,而SerpApi需要一个API Key才能使用,那么就意味着这个工具需要一个凭据才可以使用,而像`wikipedia`这样的工具,就不需要填写凭据字段,可以参考[这里](../../provider/builtin/wikipedia/wikipedia.yaml)。 + +配置好凭据字段后效果如下: +```yaml +identity: + author: Dify + name: google + label: + en_US: Google + zh_Hans: Google + description: + en_US: Google + zh_Hans: Google + icon: icon.svg +credentials_for_provider: # 凭据字段 + serpapi_api_key: # 凭据字段名称 + type: secret-input # 凭据字段类型 + required: true # 是否必填 + label: # 凭据字段标签 + en_US: SerpApi API key # 英文标签 + zh_Hans: SerpApi API key # 中文标签 + placeholder: # 凭据字段占位符 + en_US: Please input your SerpApi API key # 英文占位符 + zh_Hans: 请输入你的 SerpApi API key # 中文占位符 + help: # 凭据字段帮助文本 + en_US: Get your SerpApi API key from SerpApi # 英文帮助文本 + zh_Hans: 从 SerpApi 获取您的 SerpApi API key # 中文帮助文本 + url: https://serpapi.com/manage-api-key # 凭据字段帮助链接 + +``` + +- `type`:凭据字段类型,目前支持`secret-input`、`text-input`、`select` 三种类型,分别对应密码输入框、文本输入框、下拉框,如果为`secret-input`,则会在前端隐藏输入内容,并且后端会对输入内容进行加密。 + +## 3. 准备工具yaml +一个供应商底下可以有多个工具,每个工具都需要一个yaml文件来描述,这个文件包含了工具的基本信息、参数、输出等。 + +仍然以GoogleSearch为例,我们需要在`google`模块下创建一个`tools`模块,并创建`tools/google_search.yaml`,内容如下。 + +```yaml +identity: # 工具的基本信息 + name: google_search # 工具名称,唯一,不允许和其他工具重名 + author: Dify # 作者 + label: # 标签,用于前端展示 + en_US: GoogleSearch # 英文标签 + zh_Hans: 谷歌搜索 # 中文标签 +description: # 描述,用于前端展示 + human: # 用于前端展示的介绍,支持多语言 + en_US: A tool for performing a Google SERP search and extracting snippets and webpages.Input should be a search query. + zh_Hans: 一个用于执行 Google SERP 搜索并提取片段和网页的工具。输入应该是一个搜索查询。 + llm: A tool for performing a Google SERP search and extracting snippets and webpages.Input should be a search query. # 传递给LLM的介绍,为了使得LLM更好理解这个工具,我们建议在这里写上关于这个工具尽可能详细的信息,让LLM能够理解并使用这个工具 +parameters: # 参数列表 + - name: query # 参数名称 + type: string # 参数类型 + required: true # 是否必填 + label: # 参数标签 + en_US: Query string # 英文标签 + zh_Hans: 查询语句 # 中文标签 + human_description: # 用于前端展示的介绍,支持多语言 + en_US: used for searching + zh_Hans: 用于搜索网页内容 + llm_description: key words for searching # 传递给LLM的介绍,同上,为了使得LLM更好理解这个参数,我们建议在这里写上关于这个参数尽可能详细的信息,让LLM能够理解这个参数 + form: llm # 表单类型,llm表示这个参数需要由Agent自行推理出来,前端将不会展示这个参数 + - name: result_type + type: select # 参数类型 + required: true + options: # 下拉框选项 + - value: text + label: + en_US: text + zh_Hans: 文本 + - value: link + label: + en_US: link + zh_Hans: 链接 + default: link + label: + en_US: Result type + zh_Hans: 结果类型 + human_description: + en_US: used for selecting the result type, text or link + zh_Hans: 用于选择结果类型,使用文本还是链接进行展示 + form: form # 表单类型,form表示这个参数需要由用户在对话开始前在前端填写 + +``` + +- `identity` 字段是必须的,它包含了工具的基本信息,包括名称、作者、标签、描述等 +- `parameters` 参数列表 + - `name` (必填)参数名称,唯一,不允许和其他参数重名 + - `type` (必填)参数类型,目前支持`string`、`number`、`boolean`、`select`、`secret-input` 五种类型,分别对应字符串、数字、布尔值、下拉框、加密输入框,对于敏感信息,我们建议使用`secret-input`类型 + - `label`(必填)参数标签,用于前端展示 + - `form` (必填)表单类型,目前支持`llm`、`form`两种类型 + - 在Agent应用中,`llm`表示该参数LLM自行推理,`form`表示要使用该工具可提前设定的参数 + - 在workflow应用中,`llm`和`form`均需要前端填写,但`llm`的参数会做为工具节点的输入变量 + - `required` 是否必填 + - 在`llm`模式下,如果参数为必填,则会要求Agent必须要推理出这个参数 + - 在`form`模式下,如果参数为必填,则会要求用户在对话开始前在前端填写这个参数 + - `options` 参数选项 + - 在`llm`模式下,Dify会将所有选项传递给LLM,LLM可以根据这些选项进行推理 + - 在`form`模式下,`type`为`select`时,前端会展示这些选项 + - `default` 默认值 + - `min` 最小值,当参数类型为`number`时可以设定 + - `max` 最大值,当参数类型为`number`时可以设定 + - `human_description` 用于前端展示的介绍,支持多语言 + - `placeholder` 字段输入框的提示文字,在表单类型为`form`,参数类型为`string`、`number`、`secret-input`时,可以设定,支持多语言 + - `llm_description` 传递给LLM的介绍,为了使得LLM更好理解这个参数,我们建议在这里写上关于这个参数尽可能详细的信息,让LLM能够理解这个参数 + + +## 4. 准备工具代码 +当完成工具的配置以后,我们就可以开始编写工具代码了,主要用于实现工具的逻辑。 + +在`google/tools`模块下创建`google_search.py`,内容如下。 + +```python +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.entities.tool_entities import ToolInvokeMessage + +from typing import Any, Dict, List, Union + +class GoogleSearchTool(BuiltinTool): + def _invoke(self, + user_id: str, + tool_parameters: Dict[str, Any], + ) -> Union[ToolInvokeMessage, List[ToolInvokeMessage]]: + """ + invoke tools + """ + query = tool_parameters['query'] + result_type = tool_parameters['result_type'] + api_key = self.runtime.credentials['serpapi_api_key'] + result = SerpAPI(api_key).run(query, result_type=result_type) + + if result_type == 'text': + return self.create_text_message(text=result) + return self.create_link_message(link=result) +``` + +### 参数 +工具的整体逻辑都在`_invoke`方法中,这个方法接收两个参数:`user_id`和`tool_parameters`,分别表示用户ID和工具参数 + +### 返回数据 +在工具返回时,你可以选择返回一条消息或者多个消息,这里我们返回一条消息,使用`create_text_message`和`create_link_message`可以创建一条文本消息或者一条链接消息。如需返回多条消息,可以使用列表构建,例如`[self.create_text_message('msg1'), self.create_text_message('msg2')]` + +## 5. 准备供应商代码 +最后,我们需要在供应商模块下创建一个供应商类,用于实现供应商的凭据验证逻辑,如果凭据验证失败,将会抛出`ToolProviderCredentialValidationError`异常。 + +在`google`模块下创建`google.py`,内容如下。 + +```python +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController +from core.tools.errors import ToolProviderCredentialValidationError + +from core.tools.provider.builtin.google.tools.google_search import GoogleSearchTool + +from typing import Any, Dict + +class GoogleProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: Dict[str, Any]) -> None: + try: + # 1. 此处需要使用GoogleSearchTool()实例化一个GoogleSearchTool,它会自动加载GoogleSearchTool的yaml配置,但是此时它内部没有凭据信息 + # 2. 随后需要使用fork_tool_runtime方法,将当前的凭据信息传递给GoogleSearchTool + # 3. 最后invoke即可,参数需要根据GoogleSearchTool的yaml中配置的参数规则进行传递 + GoogleSearchTool().fork_tool_runtime( + meta={ + "credentials": credentials, + } + ).invoke( + user_id='', + tool_parameters={ + "query": "test", + "result_type": "link" + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) +``` + +## 完成 +当上述步骤完成以后,我们就可以在前端看到这个工具了,并且可以在Agent中使用这个工具。 + +当然,因为google_search需要一个凭据,在使用之前,还需要在前端配置它的凭据。 + +![Alt text](../images/index/image-2.png) diff --git a/api/core/tools/entities/api_entities.py b/api/core/tools/entities/api_entities.py new file mode 100644 index 0000000000000000000000000000000000000000..975c374cae8356a62358205ee24bfbcfa6e77592 --- /dev/null +++ b/api/core/tools/entities/api_entities.py @@ -0,0 +1,71 @@ +from typing import Literal, Optional + +from pydantic import BaseModel, Field, field_validator + +from core.model_runtime.utils.encoders import jsonable_encoder +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_entities import ToolProviderCredentials, ToolProviderType +from core.tools.tool.tool import ToolParameter + + +class UserTool(BaseModel): + author: str + name: str # identifier + label: I18nObject # label + description: I18nObject + parameters: Optional[list[ToolParameter]] = None + labels: list[str] | None = None + + +UserToolProviderTypeLiteral = Optional[Literal["builtin", "api", "workflow"]] + + +class UserToolProvider(BaseModel): + id: str + author: str + name: str # identifier + description: I18nObject + icon: str + label: I18nObject # label + type: ToolProviderType + masked_credentials: Optional[dict] = None + original_credentials: Optional[dict] = None + is_team_authorization: bool = False + allow_delete: bool = True + tools: list[UserTool] = Field(default_factory=list) + labels: list[str] | None = None + + @field_validator("tools", mode="before") + @classmethod + def convert_none_to_empty_list(cls, v): + return v if v is not None else [] + + def to_dict(self) -> dict: + # ------------- + # overwrite tool parameter types for temp fix + tools = jsonable_encoder(self.tools) + for tool in tools: + if tool.get("parameters"): + for parameter in tool.get("parameters"): + if parameter.get("type") == ToolParameter.ToolParameterType.SYSTEM_FILES.value: + parameter["type"] = "files" + # ------------- + + return { + "id": self.id, + "author": self.author, + "name": self.name, + "description": self.description.to_dict(), + "icon": self.icon, + "label": self.label.to_dict(), + "type": self.type.value, + "team_credentials": self.masked_credentials, + "is_team_authorization": self.is_team_authorization, + "allow_delete": self.allow_delete, + "tools": tools, + "labels": self.labels, + } + + +class UserToolProviderCredentials(BaseModel): + credentials: dict[str, ToolProviderCredentials] diff --git a/api/core/tools/entities/common_entities.py b/api/core/tools/entities/common_entities.py new file mode 100644 index 0000000000000000000000000000000000000000..924e6fc0cf9f1758f47cfe727c3dff2516eaf388 --- /dev/null +++ b/api/core/tools/entities/common_entities.py @@ -0,0 +1,23 @@ +from typing import Optional + +from pydantic import BaseModel, Field + + +class I18nObject(BaseModel): + """ + Model class for i18n object. + """ + + en_US: str + zh_Hans: Optional[str] = Field(default=None) + pt_BR: Optional[str] = Field(default=None) + ja_JP: Optional[str] = Field(default=None) + + def __init__(self, **data): + super().__init__(**data) + self.zh_Hans = self.zh_Hans or self.en_US + self.pt_BR = self.pt_BR or self.en_US + self.ja_JP = self.ja_JP or self.en_US + + def to_dict(self) -> dict: + return {"zh_Hans": self.zh_Hans, "en_US": self.en_US, "pt_BR": self.pt_BR, "ja_JP": self.ja_JP} diff --git a/api/core/tools/entities/tool_bundle.py b/api/core/tools/entities/tool_bundle.py new file mode 100644 index 0000000000000000000000000000000000000000..7c365dc69d3b39380d05cbc22d481f60b4f023a1 --- /dev/null +++ b/api/core/tools/entities/tool_bundle.py @@ -0,0 +1,29 @@ +from typing import Optional + +from pydantic import BaseModel + +from core.tools.entities.tool_entities import ToolParameter + + +class ApiToolBundle(BaseModel): + """ + This class is used to store the schema information of an api based tool. + such as the url, the method, the parameters, etc. + """ + + # server_url + server_url: str + # method + method: str + # summary + summary: Optional[str] = None + # operation_id + operation_id: str | None = None + # parameters + parameters: Optional[list[ToolParameter]] = None + # author + author: str + # icon + icon: Optional[str] = None + # openapi operation + openapi: dict diff --git a/api/core/tools/entities/tool_entities.py b/api/core/tools/entities/tool_entities.py new file mode 100644 index 0000000000000000000000000000000000000000..c87a90c03a6f7e102d3458d6ab971b99a6d9ae0e --- /dev/null +++ b/api/core/tools/entities/tool_entities.py @@ -0,0 +1,528 @@ +from enum import Enum, StrEnum +from typing import Any, Optional, Union, cast + +from pydantic import BaseModel, Field, field_validator + +from core.tools.entities.common_entities import I18nObject + + +class ToolLabelEnum(Enum): + SEARCH = "search" + IMAGE = "image" + VIDEOS = "videos" + WEATHER = "weather" + FINANCE = "finance" + DESIGN = "design" + TRAVEL = "travel" + SOCIAL = "social" + NEWS = "news" + MEDICAL = "medical" + PRODUCTIVITY = "productivity" + EDUCATION = "education" + BUSINESS = "business" + ENTERTAINMENT = "entertainment" + UTILITIES = "utilities" + OTHER = "other" + + +class ToolProviderType(Enum): + """ + Enum class for tool provider + """ + + BUILT_IN = "builtin" + WORKFLOW = "workflow" + API = "api" + APP = "app" + DATASET_RETRIEVAL = "dataset-retrieval" + + @classmethod + def value_of(cls, value: str) -> "ToolProviderType": + """ + Get value of given mode. + + :param value: mode value + :return: mode + """ + for mode in cls: + if mode.value == value: + return mode + raise ValueError(f"invalid mode value {value}") + + +class ApiProviderSchemaType(Enum): + """ + Enum class for api provider schema type. + """ + + OPENAPI = "openapi" + SWAGGER = "swagger" + OPENAI_PLUGIN = "openai_plugin" + OPENAI_ACTIONS = "openai_actions" + + @classmethod + def value_of(cls, value: str) -> "ApiProviderSchemaType": + """ + Get value of given mode. + + :param value: mode value + :return: mode + """ + for mode in cls: + if mode.value == value: + return mode + raise ValueError(f"invalid mode value {value}") + + +class ApiProviderAuthType(Enum): + """ + Enum class for api provider auth type. + """ + + NONE = "none" + API_KEY = "api_key" + + @classmethod + def value_of(cls, value: str) -> "ApiProviderAuthType": + """ + Get value of given mode. + + :param value: mode value + :return: mode + """ + for mode in cls: + if mode.value == value: + return mode + raise ValueError(f"invalid mode value {value}") + + +class ToolInvokeMessage(BaseModel): + class MessageType(Enum): + TEXT = "text" + IMAGE = "image" + LINK = "link" + BLOB = "blob" + JSON = "json" + IMAGE_LINK = "image_link" + FILE = "file" + + type: MessageType = MessageType.TEXT + """ + plain text, image url or link url + """ + message: str | bytes | dict | None = None + # TODO: Use a BaseModel for meta + meta: dict[str, Any] = Field(default_factory=dict) + save_as: str = "" + + +class ToolInvokeMessageBinary(BaseModel): + mimetype: str = Field(..., description="The mimetype of the binary") + url: str = Field(..., description="The url of the binary") + save_as: str = "" + file_var: Optional[dict[str, Any]] = None + + +class ToolParameterOption(BaseModel): + value: str = Field(..., description="The value of the option") + label: I18nObject = Field(..., description="The label of the option") + + @field_validator("value", mode="before") + @classmethod + def transform_id_to_str(cls, value) -> str: + if not isinstance(value, str): + return str(value) + else: + return value + + +class ToolParameter(BaseModel): + class ToolParameterType(StrEnum): + STRING = "string" + NUMBER = "number" + BOOLEAN = "boolean" + SELECT = "select" + SECRET_INPUT = "secret-input" + FILE = "file" + FILES = "files" + + # deprecated, should not use. + SYSTEM_FILES = "systme-files" + + def as_normal_type(self): + if self in { + ToolParameter.ToolParameterType.SECRET_INPUT, + ToolParameter.ToolParameterType.SELECT, + }: + return "string" + return self.value + + def cast_value(self, value: Any, /): + try: + match self: + case ( + ToolParameter.ToolParameterType.STRING + | ToolParameter.ToolParameterType.SECRET_INPUT + | ToolParameter.ToolParameterType.SELECT + ): + if value is None: + return "" + else: + return value if isinstance(value, str) else str(value) + + case ToolParameter.ToolParameterType.BOOLEAN: + if value is None: + return False + elif isinstance(value, str): + # Allowed YAML boolean value strings: https://yaml.org/type/bool.html + # and also '0' for False and '1' for True + match value.lower(): + case "true" | "yes" | "y" | "1": + return True + case "false" | "no" | "n" | "0": + return False + case _: + return bool(value) + else: + return value if isinstance(value, bool) else bool(value) + + case ToolParameter.ToolParameterType.NUMBER: + if isinstance(value, int | float): + return value + elif isinstance(value, str) and value: + if "." in value: + return float(value) + else: + return int(value) + case ( + ToolParameter.ToolParameterType.SYSTEM_FILES + | ToolParameter.ToolParameterType.FILE + | ToolParameter.ToolParameterType.FILES + ): + return value + case _: + return str(value) + + except Exception: + raise ValueError(f"The tool parameter value {value} is not in correct type.") + + class ToolParameterForm(Enum): + SCHEMA = "schema" # should be set while adding tool + FORM = "form" # should be set before invoking tool + LLM = "llm" # will be set by LLM + + name: str = Field(..., description="The name of the parameter") + label: I18nObject = Field(..., description="The label presented to the user") + human_description: Optional[I18nObject] = Field(None, description="The description presented to the user") + placeholder: Optional[I18nObject] = Field(None, description="The placeholder presented to the user") + type: ToolParameterType = Field(..., description="The type of the parameter") + form: ToolParameterForm = Field(..., description="The form of the parameter, schema/form/llm") + llm_description: Optional[str] = None + required: Optional[bool] = False + default: Optional[Union[float, int, str]] = None + min: Optional[Union[float, int]] = None + max: Optional[Union[float, int]] = None + options: Optional[list[ToolParameterOption]] = None + + @classmethod + def get_simple_instance( + cls, + name: str, + llm_description: str, + type: ToolParameterType, + required: bool, + options: Optional[list[str]] = None, + ) -> "ToolParameter": + """ + get a simple tool parameter + + :param name: the name of the parameter + :param llm_description: the description presented to the LLM + :param type: the type of the parameter + :param required: if the parameter is required + :param options: the options of the parameter + """ + # convert options to ToolParameterOption + # FIXME fix the type error + if options: + options = [ + ToolParameterOption(value=option, label=I18nObject(en_US=option, zh_Hans=option)) # type: ignore + for option in options # type: ignore + ] + return cls( + name=name, + label=I18nObject(en_US="", zh_Hans=""), + human_description=I18nObject(en_US="", zh_Hans=""), + placeholder=None, + type=type, + form=cls.ToolParameterForm.LLM, + llm_description=llm_description, + required=required, + options=options, # type: ignore + ) + + +class ToolProviderIdentity(BaseModel): + author: str = Field(..., description="The author of the tool") + name: str = Field(..., description="The name of the tool") + description: I18nObject = Field(..., description="The description of the tool") + icon: str = Field(..., description="The icon of the tool") + label: I18nObject = Field(..., description="The label of the tool") + tags: Optional[list[ToolLabelEnum]] = Field( + default=[], + description="The tags of the tool", + ) + + +class ToolDescription(BaseModel): + human: I18nObject = Field(..., description="The description presented to the user") + llm: str = Field(..., description="The description presented to the LLM") + + +class ToolIdentity(BaseModel): + author: str = Field(..., description="The author of the tool") + name: str = Field(..., description="The name of the tool") + label: I18nObject = Field(..., description="The label of the tool") + provider: str = Field(..., description="The provider of the tool") + icon: Optional[str] = None + + +class ToolCredentialsOption(BaseModel): + value: str = Field(..., description="The value of the option") + label: I18nObject = Field(..., description="The label of the option") + + +class ToolProviderCredentials(BaseModel): + class CredentialsType(Enum): + SECRET_INPUT = "secret-input" + TEXT_INPUT = "text-input" + SELECT = "select" + BOOLEAN = "boolean" + + @classmethod + def value_of(cls, value: str) -> "ToolProviderCredentials.CredentialsType": + """ + Get value of given mode. + + :param value: mode value + :return: mode + """ + for mode in cls: + if mode.value == value: + return mode + raise ValueError(f"invalid mode value {value}") + + @staticmethod + def default(value: str) -> str: + return "" + + name: str = Field(..., description="The name of the credentials") + type: CredentialsType = Field(..., description="The type of the credentials") + required: bool = False + default: Optional[Union[int, str]] = None + options: Optional[list[ToolCredentialsOption]] = None + label: Optional[I18nObject] = None + help: Optional[I18nObject] = None + url: Optional[str] = None + placeholder: Optional[I18nObject] = None + + def to_dict(self) -> dict: + return { + "name": self.name, + "type": self.type.value, + "required": self.required, + "default": self.default, + "options": self.options, + "help": self.help.to_dict() if self.help else None, + "label": self.label.to_dict() if self.label else None, + "url": self.url, + "placeholder": self.placeholder.to_dict() if self.placeholder else None, + } + + +class ToolRuntimeVariableType(Enum): + TEXT = "text" + IMAGE = "image" + + +class ToolRuntimeVariable(BaseModel): + type: ToolRuntimeVariableType = Field(..., description="The type of the variable") + name: str = Field(..., description="The name of the variable") + position: int = Field(..., description="The position of the variable") + tool_name: str = Field(..., description="The name of the tool") + + +class ToolRuntimeTextVariable(ToolRuntimeVariable): + value: str = Field(..., description="The value of the variable") + + +class ToolRuntimeImageVariable(ToolRuntimeVariable): + value: str = Field(..., description="The path of the image") + + +class ToolRuntimeVariablePool(BaseModel): + conversation_id: str = Field(..., description="The conversation id") + user_id: str = Field(..., description="The user id") + tenant_id: str = Field(..., description="The tenant id of assistant") + + pool: list[ToolRuntimeVariable] = Field(..., description="The pool of variables") + + def __init__(self, **data: Any): + pool = data.get("pool", []) + # convert pool into correct type + for index, variable in enumerate(pool): + if variable["type"] == ToolRuntimeVariableType.TEXT.value: + pool[index] = ToolRuntimeTextVariable(**variable) + elif variable["type"] == ToolRuntimeVariableType.IMAGE.value: + pool[index] = ToolRuntimeImageVariable(**variable) + super().__init__(**data) + + def dict(self) -> dict: # type: ignore + """ + FIXME: just ignore the type check for now + """ + return { + "conversation_id": self.conversation_id, + "user_id": self.user_id, + "tenant_id": self.tenant_id, + "pool": [variable.model_dump() for variable in self.pool], + } + + def set_text(self, tool_name: str, name: str, value: str) -> None: + """ + set a text variable + """ + for variable in self.pool: + if variable.name == name: + if variable.type == ToolRuntimeVariableType.TEXT: + variable = cast(ToolRuntimeTextVariable, variable) + variable.value = value + return + + variable = ToolRuntimeTextVariable( + type=ToolRuntimeVariableType.TEXT, + name=name, + position=len(self.pool), + tool_name=tool_name, + value=value, + ) + + self.pool.append(variable) + + def set_file(self, tool_name: str, value: str, name: Optional[str] = None) -> None: + """ + set an image variable + + :param tool_name: the name of the tool + :param value: the id of the file + """ + # check how many image variables are there + image_variable_count = 0 + for variable in self.pool: + if variable.type == ToolRuntimeVariableType.IMAGE: + image_variable_count += 1 + + if name is None: + name = f"file_{image_variable_count}" + + for variable in self.pool: + if variable.name == name: + if variable.type == ToolRuntimeVariableType.IMAGE: + variable = cast(ToolRuntimeImageVariable, variable) + variable.value = value + return + + variable = ToolRuntimeImageVariable( + type=ToolRuntimeVariableType.IMAGE, + name=name, + position=len(self.pool), + tool_name=tool_name, + value=value, + ) + + self.pool.append(variable) + + +class ModelToolPropertyKey(Enum): + IMAGE_PARAMETER_NAME = "image_parameter_name" + + +class ModelToolConfiguration(BaseModel): + """ + Model tool configuration + """ + + type: str = Field(..., description="The type of the model tool") + model: str = Field(..., description="The model") + label: I18nObject = Field(..., description="The label of the model tool") + properties: dict[ModelToolPropertyKey, Any] = Field(..., description="The properties of the model tool") + + +class ModelToolProviderConfiguration(BaseModel): + """ + Model tool provider configuration + """ + + provider: str = Field(..., description="The provider of the model tool") + models: list[ModelToolConfiguration] = Field(..., description="The models of the model tool") + label: I18nObject = Field(..., description="The label of the model tool") + + +class WorkflowToolParameterConfiguration(BaseModel): + """ + Workflow tool configuration + """ + + name: str = Field(..., description="The name of the parameter") + description: str = Field(..., description="The description of the parameter") + form: ToolParameter.ToolParameterForm = Field(..., description="The form of the parameter") + + +class ToolInvokeMeta(BaseModel): + """ + Tool invoke meta + """ + + time_cost: float = Field(..., description="The time cost of the tool invoke") + error: Optional[str] = None + tool_config: Optional[dict] = None + + @classmethod + def empty(cls) -> "ToolInvokeMeta": + """ + Get an empty instance of ToolInvokeMeta + """ + return cls(time_cost=0.0, error=None, tool_config={}) + + @classmethod + def error_instance(cls, error: str) -> "ToolInvokeMeta": + """ + Get an instance of ToolInvokeMeta with error + """ + return cls(time_cost=0.0, error=error, tool_config={}) + + def to_dict(self) -> dict: + return { + "time_cost": self.time_cost, + "error": self.error, + "tool_config": self.tool_config, + } + + +class ToolLabel(BaseModel): + """ + Tool label + """ + + name: str = Field(..., description="The name of the tool") + label: I18nObject = Field(..., description="The label of the tool") + icon: str = Field(..., description="The icon of the tool") + + +class ToolInvokeFrom(Enum): + """ + Enum class for tool invoke + """ + + WORKFLOW = "workflow" + AGENT = "agent" diff --git a/api/core/tools/entities/values.py b/api/core/tools/entities/values.py new file mode 100644 index 0000000000000000000000000000000000000000..f460df7e25c916b702b0bda59abd9c42e20e44e3 --- /dev/null +++ b/api/core/tools/entities/values.py @@ -0,0 +1,111 @@ +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_entities import ToolLabel, ToolLabelEnum + +ICONS = { + ToolLabelEnum.SEARCH: """ + +""", # noqa: E501 + ToolLabelEnum.IMAGE: """ + +""", # noqa: E501 + ToolLabelEnum.VIDEOS: """ + +""", # noqa: E501 + ToolLabelEnum.WEATHER: """ + +""", # noqa: E501 + ToolLabelEnum.FINANCE: """ + +""", # noqa: E501 + ToolLabelEnum.DESIGN: """ + +""", # noqa: E501 + ToolLabelEnum.TRAVEL: """ + +""", # noqa: E501 + ToolLabelEnum.SOCIAL: """ + +""", # noqa: E501 + ToolLabelEnum.NEWS: """ + +""", # noqa: E501 + ToolLabelEnum.MEDICAL: """ + +""", # noqa: E501 + ToolLabelEnum.PRODUCTIVITY: """ + +""", # noqa: E501 + ToolLabelEnum.EDUCATION: """ + +""", # noqa: E501 + ToolLabelEnum.BUSINESS: """ + +""", # noqa: E501 + ToolLabelEnum.ENTERTAINMENT: """ + +""", # noqa: E501 + ToolLabelEnum.UTILITIES: """ + +""", # noqa: E501 + ToolLabelEnum.OTHER: """ + +""", # noqa: E501 +} + +default_tool_label_dict = { + ToolLabelEnum.SEARCH: ToolLabel( + name="search", label=I18nObject(en_US="Search", zh_Hans="搜索"), icon=ICONS[ToolLabelEnum.SEARCH] + ), + ToolLabelEnum.IMAGE: ToolLabel( + name="image", label=I18nObject(en_US="Image", zh_Hans="图片"), icon=ICONS[ToolLabelEnum.IMAGE] + ), + ToolLabelEnum.VIDEOS: ToolLabel( + name="videos", label=I18nObject(en_US="Videos", zh_Hans="视频"), icon=ICONS[ToolLabelEnum.VIDEOS] + ), + ToolLabelEnum.WEATHER: ToolLabel( + name="weather", label=I18nObject(en_US="Weather", zh_Hans="天气"), icon=ICONS[ToolLabelEnum.WEATHER] + ), + ToolLabelEnum.FINANCE: ToolLabel( + name="finance", label=I18nObject(en_US="Finance", zh_Hans="金融"), icon=ICONS[ToolLabelEnum.FINANCE] + ), + ToolLabelEnum.DESIGN: ToolLabel( + name="design", label=I18nObject(en_US="Design", zh_Hans="设计"), icon=ICONS[ToolLabelEnum.DESIGN] + ), + ToolLabelEnum.TRAVEL: ToolLabel( + name="travel", label=I18nObject(en_US="Travel", zh_Hans="旅行"), icon=ICONS[ToolLabelEnum.TRAVEL] + ), + ToolLabelEnum.SOCIAL: ToolLabel( + name="social", label=I18nObject(en_US="Social", zh_Hans="社交"), icon=ICONS[ToolLabelEnum.SOCIAL] + ), + ToolLabelEnum.NEWS: ToolLabel( + name="news", label=I18nObject(en_US="News", zh_Hans="新闻"), icon=ICONS[ToolLabelEnum.NEWS] + ), + ToolLabelEnum.MEDICAL: ToolLabel( + name="medical", label=I18nObject(en_US="Medical", zh_Hans="医疗"), icon=ICONS[ToolLabelEnum.MEDICAL] + ), + ToolLabelEnum.PRODUCTIVITY: ToolLabel( + name="productivity", + label=I18nObject(en_US="Productivity", zh_Hans="生产力"), + icon=ICONS[ToolLabelEnum.PRODUCTIVITY], + ), + ToolLabelEnum.EDUCATION: ToolLabel( + name="education", label=I18nObject(en_US="Education", zh_Hans="教育"), icon=ICONS[ToolLabelEnum.EDUCATION] + ), + ToolLabelEnum.BUSINESS: ToolLabel( + name="business", label=I18nObject(en_US="Business", zh_Hans="商业"), icon=ICONS[ToolLabelEnum.BUSINESS] + ), + ToolLabelEnum.ENTERTAINMENT: ToolLabel( + name="entertainment", + label=I18nObject(en_US="Entertainment", zh_Hans="娱乐"), + icon=ICONS[ToolLabelEnum.ENTERTAINMENT], + ), + ToolLabelEnum.UTILITIES: ToolLabel( + name="utilities", label=I18nObject(en_US="Utilities", zh_Hans="工具"), icon=ICONS[ToolLabelEnum.UTILITIES] + ), + ToolLabelEnum.OTHER: ToolLabel( + name="other", label=I18nObject(en_US="Other", zh_Hans="其他"), icon=ICONS[ToolLabelEnum.OTHER] + ), +} + +default_tool_labels = [v for k, v in default_tool_label_dict.items()] +default_tool_label_name_list = [label.name for label in default_tool_labels] diff --git a/api/core/tools/errors.py b/api/core/tools/errors.py new file mode 100644 index 0000000000000000000000000000000000000000..c5f9ca477401f2bca0cd6007f6e19d634db56482 --- /dev/null +++ b/api/core/tools/errors.py @@ -0,0 +1,37 @@ +from core.tools.entities.tool_entities import ToolInvokeMeta + + +class ToolProviderNotFoundError(ValueError): + pass + + +class ToolNotFoundError(ValueError): + pass + + +class ToolParameterValidationError(ValueError): + pass + + +class ToolProviderCredentialValidationError(ValueError): + pass + + +class ToolNotSupportedError(ValueError): + pass + + +class ToolInvokeError(ValueError): + pass + + +class ToolApiSchemaError(ValueError): + pass + + +class ToolEngineInvokeError(Exception): + meta: ToolInvokeMeta + + def __init__(self, meta, **kwargs): + self.meta = meta + super().__init__(**kwargs) diff --git a/api/core/tools/provider/_position.yaml b/api/core/tools/provider/_position.yaml new file mode 100644 index 0000000000000000000000000000000000000000..937fb40774cbac62f6c2d291562747705a8ae0ad --- /dev/null +++ b/api/core/tools/provider/_position.yaml @@ -0,0 +1,81 @@ +- google +- bing +- perplexity +- duckduckgo +- searchapi +- serper +- searxng +- websearch +- tavily +- stackexchange +- pubmed +- arxiv +- aws +- nominatim +- devdocs +- spider +- firecrawl +- brave +- crossref +- jina +- webscraper +- dalle +- azuredalle +- stability +- stablediffusion +- cogview +- comfyui +- getimgai +- siliconflow +- spark +- stepfun +- xinference +- alphavantage +- yahoo +- openweather +- gaode +- aippt +- chart +- youtube +- did +- dingtalk +- discord +- feishu +- feishu_base +- feishu_document +- feishu_message +- feishu_wiki +- feishu_task +- feishu_calendar +- feishu_spreadsheet +- lark_base +- lark_document +- lark_message_and_group +- lark_wiki +- lark_task +- lark_calendar +- lark_spreadsheet +- slack +- twilio +- wecom +- wikipedia +- code +- wolframalpha +- maths +- github +- gitlab +- time +- vectorizer +- qrcode +- tianditu +- aliyuque +- google_translate +- hap +- json_process +- judge0ce +- novitaai +- onebot +- regex +- trello +- vanna +- fal diff --git a/api/core/tools/provider/api_tool_provider.py b/api/core/tools/provider/api_tool_provider.py new file mode 100644 index 0000000000000000000000000000000000000000..f451edbf2ee969ebc4877d8308a8560242fa341c --- /dev/null +++ b/api/core/tools/provider/api_tool_provider.py @@ -0,0 +1,172 @@ +from typing import Optional + +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_bundle import ApiToolBundle +from core.tools.entities.tool_entities import ( + ApiProviderAuthType, + ToolCredentialsOption, + ToolDescription, + ToolIdentity, + ToolProviderCredentials, + ToolProviderIdentity, + ToolProviderType, +) +from core.tools.provider.tool_provider import ToolProviderController +from core.tools.tool.api_tool import ApiTool +from core.tools.tool.tool import Tool +from extensions.ext_database import db +from models.tools import ApiToolProvider + + +class ApiToolProviderController(ToolProviderController): + provider_id: str + + @staticmethod + def from_db(db_provider: ApiToolProvider, auth_type: ApiProviderAuthType) -> "ApiToolProviderController": + credentials_schema = { + "auth_type": ToolProviderCredentials( + name="auth_type", + required=True, + type=ToolProviderCredentials.CredentialsType.SELECT, + options=[ + ToolCredentialsOption(value="none", label=I18nObject(en_US="None", zh_Hans="无")), + ToolCredentialsOption(value="api_key", label=I18nObject(en_US="api_key", zh_Hans="api_key")), + ], + default="none", + help=I18nObject(en_US="The auth type of the api provider", zh_Hans="api provider 的认证类型"), + ) + } + if auth_type == ApiProviderAuthType.API_KEY: + credentials_schema = { + **credentials_schema, + "api_key_header": ToolProviderCredentials( + name="api_key_header", + required=False, + default="api_key", + type=ToolProviderCredentials.CredentialsType.TEXT_INPUT, + help=I18nObject(en_US="The header name of the api key", zh_Hans="携带 api key 的 header 名称"), + ), + "api_key_value": ToolProviderCredentials( + name="api_key_value", + required=True, + type=ToolProviderCredentials.CredentialsType.SECRET_INPUT, + help=I18nObject(en_US="The api key", zh_Hans="api key的值"), + ), + "api_key_header_prefix": ToolProviderCredentials( + name="api_key_header_prefix", + required=False, + default="basic", + type=ToolProviderCredentials.CredentialsType.SELECT, + help=I18nObject(en_US="The prefix of the api key header", zh_Hans="api key header 的前缀"), + options=[ + ToolCredentialsOption(value="basic", label=I18nObject(en_US="Basic", zh_Hans="Basic")), + ToolCredentialsOption(value="bearer", label=I18nObject(en_US="Bearer", zh_Hans="Bearer")), + ToolCredentialsOption(value="custom", label=I18nObject(en_US="Custom", zh_Hans="Custom")), + ], + ), + } + elif auth_type == ApiProviderAuthType.NONE: + pass + else: + raise ValueError(f"invalid auth type {auth_type}") + user_name = db_provider.user.name if db_provider.user_id and db_provider.user is not None else "" + return ApiToolProviderController( + identity=ToolProviderIdentity( + author=user_name, + name=db_provider.name, + label=I18nObject(en_US=db_provider.name, zh_Hans=db_provider.name), + description=I18nObject(en_US=db_provider.description, zh_Hans=db_provider.description), + icon=db_provider.icon, + ), + credentials_schema=credentials_schema, + provider_id=db_provider.id or "", + tools=None, + ) + + @property + def provider_type(self) -> ToolProviderType: + return ToolProviderType.API + + def _parse_tool_bundle(self, tool_bundle: ApiToolBundle) -> ApiTool: + """ + parse tool bundle to tool + + :param tool_bundle: the tool bundle + :return: the tool + """ + return ApiTool( + api_bundle=tool_bundle, + identity=ToolIdentity( + author=tool_bundle.author, + name=tool_bundle.operation_id or "", + label=I18nObject(en_US=tool_bundle.operation_id, zh_Hans=tool_bundle.operation_id), + icon=self.identity.icon if self.identity else None, + provider=self.provider_id, + ), + description=ToolDescription( + human=I18nObject(en_US=tool_bundle.summary or "", zh_Hans=tool_bundle.summary or ""), + llm=tool_bundle.summary or "", + ), + parameters=tool_bundle.parameters or [], + ) + + def load_bundled_tools(self, tools: list[ApiToolBundle]) -> list[Tool]: + """ + load bundled tools + + :param tools: the bundled tools + :return: the tools + """ + self.tools = [self._parse_tool_bundle(tool) for tool in tools] + + return self.tools + + def get_tools(self, user_id: str = "", tenant_id: str = "") -> Optional[list[Tool]]: + """ + fetch tools from database + + :param user_id: the user id + :param tenant_id: the tenant id + :return: the tools + """ + if self.tools is not None: + return self.tools + if self.identity is None: + return None + + tools: list[Tool] = [] + + # get tenant api providers + db_providers: list[ApiToolProvider] = ( + db.session.query(ApiToolProvider) + .filter(ApiToolProvider.tenant_id == tenant_id, ApiToolProvider.name == self.identity.name) + .all() + ) + + if db_providers and len(db_providers) != 0: + for db_provider in db_providers: + for tool in db_provider.tools: + assistant_tool = self._parse_tool_bundle(tool) + assistant_tool.is_team_authorization = True + tools.append(assistant_tool) + + self.tools = tools + return tools + + def get_tool(self, tool_name: str) -> Tool: + """ + get tool by name + + :param tool_name: the name of the tool + :return: the tool + """ + if self.tools is None: + self.get_tools() + + for tool in self.tools or []: + if tool.identity is None: + continue + if tool.identity.name == tool_name: + return tool + + raise ValueError(f"tool {tool_name} not found") diff --git a/api/core/tools/provider/app_tool_provider.py b/api/core/tools/provider/app_tool_provider.py new file mode 100644 index 0000000000000000000000000000000000000000..fc29920acd40dc1dc29ff175fac92ef37a21f186 --- /dev/null +++ b/api/core/tools/provider/app_tool_provider.py @@ -0,0 +1,106 @@ +import logging +from typing import Any, Optional + +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_entities import ToolParameter, ToolParameterOption, ToolProviderType +from core.tools.provider.tool_provider import ToolProviderController +from core.tools.tool.api_tool import ApiTool +from core.tools.tool.tool import Tool +from extensions.ext_database import db +from models.model import App, AppModelConfig +from models.tools import PublishedAppTool + +logger = logging.getLogger(__name__) + + +class AppToolProviderEntity(ToolProviderController): + @property + def provider_type(self) -> ToolProviderType: + return ToolProviderType.APP + + def _validate_credentials(self, tool_name: str, credentials: dict[str, Any]) -> None: + pass + + def validate_parameters(self, tool_id: int, tool_name: str, tool_parameters: dict[str, Any]) -> None: + pass + + def get_tools(self, user_id: str = "", tenant_id: str = "") -> list[Tool]: + db_tools: list[PublishedAppTool] = ( + db.session.query(PublishedAppTool) + .filter( + PublishedAppTool.user_id == user_id, + ) + .all() + ) + + if not db_tools or len(db_tools) == 0: + return [] + + tools: list[Tool] = [] + + for db_tool in db_tools: + tool: dict[str, Any] = { + "identity": { + "author": db_tool.author, + "name": db_tool.tool_name, + "label": {"en_US": db_tool.tool_name, "zh_Hans": db_tool.tool_name}, + "icon": "", + }, + "description": { + "human": {"en_US": db_tool.description_i18n.en_US, "zh_Hans": db_tool.description_i18n.zh_Hans}, + "llm": db_tool.llm_description, + }, + "parameters": [], + } + # get app from db + app: Optional[App] = db_tool.app + + if not app: + logger.error(f"app {db_tool.app_id} not found") + continue + + app_model_config: AppModelConfig = app.app_model_config + user_input_form_list = app_model_config.user_input_form_list + for input_form in user_input_form_list: + # get type + form_type = list(input_form.keys())[0] + default = input_form[form_type]["default"] + required = input_form[form_type]["required"] + label = input_form[form_type]["label"] + variable_name = input_form[form_type]["variable_name"] + options = input_form[form_type].get("options", []) + if form_type in {"paragraph", "text-input"}: + tool["parameters"].append( + ToolParameter( + name=variable_name, + label=I18nObject(en_US=label, zh_Hans=label), + human_description=I18nObject(en_US=label, zh_Hans=label), + llm_description=label, + form=ToolParameter.ToolParameterForm.FORM, + type=ToolParameter.ToolParameterType.STRING, + required=required, + default=default, + placeholder=I18nObject(en_US="", zh_Hans=""), + ) + ) + elif form_type == "select": + tool["parameters"].append( + ToolParameter( + name=variable_name, + label=I18nObject(en_US=label, zh_Hans=label), + human_description=I18nObject(en_US=label, zh_Hans=label), + llm_description=label, + form=ToolParameter.ToolParameterForm.FORM, + type=ToolParameter.ToolParameterType.SELECT, + required=required, + default=default, + placeholder=I18nObject(en_US="", zh_Hans=""), + options=[ + ToolParameterOption(value=option, label=I18nObject(en_US=option, zh_Hans=option)) + for option in options + ], + ) + ) + + tools.append(ApiTool(**tool)) + return tools diff --git a/api/core/tools/provider/builtin/__init__.py b/api/core/tools/provider/builtin/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/tools/provider/builtin/_positions.py b/api/core/tools/provider/builtin/_positions.py new file mode 100644 index 0000000000000000000000000000000000000000..99a062f8c366aa1f26a6d9c36e0ba156a88d85b5 --- /dev/null +++ b/api/core/tools/provider/builtin/_positions.py @@ -0,0 +1,20 @@ +import os.path + +from core.helper.position_helper import get_tool_position_map, sort_by_position_map +from core.tools.entities.api_entities import UserToolProvider + + +class BuiltinToolProviderSort: + _position: dict[str, int] = {} + + @classmethod + def sort(cls, providers: list[UserToolProvider]) -> list[UserToolProvider]: + if not cls._position: + cls._position = get_tool_position_map(os.path.join(os.path.dirname(__file__), "..")) + + def name_func(provider: UserToolProvider) -> str: + return provider.name + + sorted_providers = sort_by_position_map(cls._position, providers, name_func) + + return sorted_providers diff --git a/api/core/tools/provider/builtin/aippt/_assets/icon.png b/api/core/tools/provider/builtin/aippt/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..b70618b4878984f88dd1991a64cc69bc0afeaaa8 Binary files /dev/null and b/api/core/tools/provider/builtin/aippt/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/aippt/aippt.py b/api/core/tools/provider/builtin/aippt/aippt.py new file mode 100644 index 0000000000000000000000000000000000000000..e0cbbd2992a5157fb33b5c543da5e8c844ce5156 --- /dev/null +++ b/api/core/tools/provider/builtin/aippt/aippt.py @@ -0,0 +1,11 @@ +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.aippt.tools.aippt import AIPPTGenerateTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class AIPPTProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + try: + AIPPTGenerateTool._get_api_token(credentials, user_id="__dify_system__") + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/aippt/aippt.yaml b/api/core/tools/provider/builtin/aippt/aippt.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9b1b45d0f21a7316d6f9e121d7917b89f47fe2b3 --- /dev/null +++ b/api/core/tools/provider/builtin/aippt/aippt.yaml @@ -0,0 +1,45 @@ +identity: + author: Dify + name: aippt + label: + en_US: AIPPT + zh_Hans: AIPPT + description: + en_US: AI-generated PPT with one click, input your content topic, and let AI serve you one-stop + zh_Hans: AI一键生成PPT,输入你的内容主题,让AI为你一站式服务到底 + icon: icon.png + tags: + - productivity + - design +credentials_for_provider: + aippt_access_key: + type: secret-input + required: true + label: + en_US: AIPPT API key + zh_Hans: AIPPT API key + pt_BR: AIPPT API key + help: + en_US: Please input your AIPPT API key + zh_Hans: 请输入你的 AIPPT API key + pt_BR: Please input your AIPPT API key + placeholder: + en_US: Please input your AIPPT API key + zh_Hans: 请输入你的 AIPPT API key + pt_BR: Please input your AIPPT API key + url: https://www.aippt.cn + aippt_secret_key: + type: secret-input + required: true + label: + en_US: AIPPT Secret key + zh_Hans: AIPPT Secret key + pt_BR: AIPPT Secret key + help: + en_US: Please input your AIPPT Secret key + zh_Hans: 请输入你的 AIPPT Secret key + pt_BR: Please input your AIPPT Secret key + placeholder: + en_US: Please input your AIPPT Secret key + zh_Hans: 请输入你的 AIPPT Secret key + pt_BR: Please input your AIPPT Secret key diff --git a/api/core/tools/provider/builtin/aippt/tools/aippt.py b/api/core/tools/provider/builtin/aippt/tools/aippt.py new file mode 100644 index 0000000000000000000000000000000000000000..0430a6654ccfac58dadf1c10d4b9f55d480bab2e --- /dev/null +++ b/api/core/tools/provider/builtin/aippt/tools/aippt.py @@ -0,0 +1,524 @@ +from base64 import b64encode +from hashlib import sha1 +from hmac import new as hmac_new +from json import loads as json_loads +from threading import Lock +from time import sleep, time +from typing import Any, Union + +from httpx import get, post +from requests import get as requests_get +from yarl import URL + +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter, ToolParameterOption +from core.tools.tool.builtin_tool import BuiltinTool + + +class AIPPTGenerateToolAdapter: + """ + A tool for generating a ppt + """ + + _api_base_url = URL("https://co.aippt.cn/api") + _api_token_cache: dict[str, dict[str, Union[str, float]]] = {} + _style_cache: dict[str, dict[str, Union[list[dict[str, Any]], float]]] = {} + + _api_token_cache_lock: Lock = Lock() + _style_cache_lock: Lock = Lock() + + _task: dict[str, Any] = {} + _task_type_map = { + "auto": 1, + "markdown": 7, + } + _tool: BuiltinTool | None + + def __init__(self, tool: BuiltinTool | None = None): + self._tool = tool + + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Invokes the AIPPT generate tool with the given user ID and tool parameters. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (dict[str, Any]): The parameters for the tool + + Returns: + ToolInvokeMessage | list[ToolInvokeMessage]: The result of the tool invocation, + which can be a single message or a list of messages. + """ + title = tool_parameters.get("title", "") + if not title: + return self._tool.create_text_message("Please provide a title for the ppt") + + model = tool_parameters.get("model", "aippt") + if not model: + return self._tool.create_text_message("Please provide a model for the ppt") + + outline = tool_parameters.get("outline", "") + + # create task + task_id = self._create_task( + type=self._task_type_map["auto" if not outline else "markdown"], + title=title, + content=outline, + user_id=user_id, + ) + + # get suit + color: str = tool_parameters.get("color", "") + style: str = tool_parameters.get("style", "") + + if color == "__default__": + color_id = "" + else: + color_id = int(color.split("-")[1]) + + if style == "__default__": + style_id = "" + else: + style_id = int(style.split("-")[1]) + + suit_id = self._get_suit(style_id=style_id, colour_id=color_id) + + # generate outline + if not outline: + self._generate_outline(task_id=task_id, model=model, user_id=user_id) + + # generate content + self._generate_content(task_id=task_id, model=model, user_id=user_id) + + # generate ppt + _, ppt_url = self._generate_ppt(task_id=task_id, suit_id=suit_id, user_id=user_id) + + return self._tool.create_text_message( + """the ppt has been created successfully,""" + f"""the ppt url is {ppt_url} .""" + """please give the ppt url to user and direct user to download it.""" + ) + + def _create_task(self, type: int, title: str, content: str, user_id: str) -> str: + """ + Create a task + + :param type: the task type + :param title: the task title + :param content: the task content + + :return: the task ID + """ + headers = { + "x-channel": "", + "x-api-key": self._tool.runtime.credentials["aippt_access_key"], + "x-token": self._get_api_token(credentials=self._tool.runtime.credentials, user_id=user_id), + } + response = post( + str(self._api_base_url / "ai" / "chat" / "v2" / "task"), + headers=headers, + files={"type": ("", str(type)), "title": ("", title), "content": ("", content)}, + ) + + if response.status_code != 200: + raise Exception(f"Failed to connect to aippt: {response.text}") + + response = response.json() + if response.get("code") != 0: + raise Exception(f"Failed to create task: {response.get('msg')}") + + return response.get("data", {}).get("id") + + def _generate_outline(self, task_id: str, model: str, user_id: str) -> str: + api_url = ( + self._api_base_url / "ai" / "chat" / "outline" + if model == "aippt" + else self._api_base_url / "ai" / "chat" / "wx" / "outline" + ) + api_url %= {"task_id": task_id} + + headers = { + "x-channel": "", + "x-api-key": self._tool.runtime.credentials["aippt_access_key"], + "x-token": self._get_api_token(credentials=self._tool.runtime.credentials, user_id=user_id), + } + + response = requests_get(url=api_url, headers=headers, stream=True, timeout=(10, 60)) + + if response.status_code != 200: + raise Exception(f"Failed to connect to aippt: {response.text}") + + outline = "" + for chunk in response.iter_lines(delimiter=b"\n\n"): + if not chunk: + continue + + event = "" + lines = chunk.decode("utf-8").split("\n") + for line in lines: + if line.startswith("event:"): + event = line[6:] + elif line.startswith("data:"): + data = line[5:] + if event == "message": + try: + data = json_loads(data) + outline += data.get("content", "") + except Exception as e: + pass + elif event == "close": + break + elif event in {"error", "filter"}: + raise Exception(f"Failed to generate outline: {data}") + + return outline + + def _generate_content(self, task_id: str, model: str, user_id: str) -> str: + api_url = ( + self._api_base_url / "ai" / "chat" / "content" + if model == "aippt" + else self._api_base_url / "ai" / "chat" / "wx" / "content" + ) + api_url %= {"task_id": task_id} + + headers = { + "x-channel": "", + "x-api-key": self._tool.runtime.credentials["aippt_access_key"], + "x-token": self._get_api_token(credentials=self._tool.runtime.credentials, user_id=user_id), + } + + response = requests_get(url=api_url, headers=headers, stream=True, timeout=(10, 60)) + + if response.status_code != 200: + raise Exception(f"Failed to connect to aippt: {response.text}") + + if model == "aippt": + content = "" + for chunk in response.iter_lines(delimiter=b"\n\n"): + if not chunk: + continue + + event = "" + lines = chunk.decode("utf-8").split("\n") + for line in lines: + if line.startswith("event:"): + event = line[6:] + elif line.startswith("data:"): + data = line[5:] + if event == "message": + try: + data = json_loads(data) + content += data.get("content", "") + except Exception as e: + pass + elif event == "close": + break + elif event in {"error", "filter"}: + raise Exception(f"Failed to generate content: {data}") + + return content + elif model == "wenxin": + response = response.json() + if response.get("code") != 0: + raise Exception(f"Failed to generate content: {response.get('msg')}") + + return response.get("data", "") + + return "" + + def _generate_ppt(self, task_id: str, suit_id: int, user_id: str) -> tuple[str, str]: + """ + Generate a ppt + + :param task_id: the task ID + :param suit_id: the suit ID + :return: the cover url of the ppt and the ppt url + """ + headers = { + "x-channel": "", + "x-api-key": self._tool.runtime.credentials["aippt_access_key"], + "x-token": self._get_api_token(credentials=self._tool.runtime.credentials, user_id=user_id), + } + + response = post( + str(self._api_base_url / "design" / "v2" / "save"), + headers=headers, + data={"task_id": task_id, "template_id": suit_id}, + timeout=(10, 60), + ) + + if response.status_code != 200: + raise Exception(f"Failed to connect to aippt: {response.text}") + + response = response.json() + if response.get("code") != 0: + raise Exception(f"Failed to generate ppt: {response.get('msg')}") + + id = response.get("data", {}).get("id") + cover_url = response.get("data", {}).get("cover_url") + + response = post( + str(self._api_base_url / "download" / "export" / "file"), + headers=headers, + data={"id": id, "format": "ppt", "files_to_zip": False, "edit": True}, + ) + + if response.status_code != 200: + raise Exception(f"Failed to connect to aippt: {response.text}") + + response = response.json() + if response.get("code") != 0: + raise Exception(f"Failed to generate ppt: {response.get('msg')}") + + export_code = response.get("data") + if not export_code: + raise Exception("Failed to generate ppt, the export code is empty") + + current_iteration = 0 + while current_iteration < 50: + # get ppt url + response = post( + str(self._api_base_url / "download" / "export" / "file" / "result"), + headers=headers, + data={"task_key": export_code}, + ) + + if response.status_code != 200: + raise Exception(f"Failed to connect to aippt: {response.text}") + + response = response.json() + if response.get("code") != 0: + raise Exception(f"Failed to generate ppt: {response.get('msg')}") + + if response.get("msg") == "导出中": + current_iteration += 1 + sleep(2) + continue + + ppt_url = response.get("data", []) + if len(ppt_url) == 0: + raise Exception("Failed to generate ppt, the ppt url is empty") + + return cover_url, ppt_url[0] + + raise Exception("Failed to generate ppt, the export is timeout") + + @classmethod + def _get_api_token(cls, credentials: dict[str, str], user_id: str) -> str: + """ + Get API token + + :param credentials: the credentials + :return: the API token + """ + access_key = credentials["aippt_access_key"] + secret_key = credentials["aippt_secret_key"] + + cache_key = f"{access_key}#@#{user_id}" + + with cls._api_token_cache_lock: + # clear expired tokens + now = time() + for key in list(cls._api_token_cache.keys()): + if cls._api_token_cache[key]["expire"] < now: + del cls._api_token_cache[key] + + if cache_key in cls._api_token_cache: + return cls._api_token_cache[cache_key]["token"] + + # get token + headers = { + "x-api-key": access_key, + "x-timestamp": str(int(now)), + "x-signature": cls._calculate_sign(access_key, secret_key, int(now)), + } + + param = {"uid": user_id, "channel": ""} + + response = get(str(cls._api_base_url / "grant" / "token"), params=param, headers=headers) + + if response.status_code != 200: + raise Exception(f"Failed to connect to aippt: {response.text}") + response = response.json() + if response.get("code") != 0: + raise Exception(f"Failed to connect to aippt: {response.get('msg')}") + + token = response.get("data", {}).get("token") + expire = response.get("data", {}).get("time_expire") + + with cls._api_token_cache_lock: + cls._api_token_cache[cache_key] = {"token": token, "expire": now + expire} + + return token + + @staticmethod + def _calculate_sign(access_key: str, secret_key: str, timestamp: int) -> str: + return b64encode( + hmac_new( + key=secret_key.encode("utf-8"), + msg=f"GET@/api/grant/token/@{timestamp}".encode(), + digestmod=sha1, + ).digest() + ).decode("utf-8") + + @classmethod + def _get_styles( + cls, credentials: dict[str, str], user_id: str + ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + """ + Get styles + """ + + # check cache + with cls._style_cache_lock: + # clear expired styles + now = time() + for key in list(cls._style_cache.keys()): + if cls._style_cache[key]["expire"] < now: + del cls._style_cache[key] + + key = f"{credentials['aippt_access_key']}#@#{user_id}" + if key in cls._style_cache: + return cls._style_cache[key]["colors"], cls._style_cache[key]["styles"] + + headers = { + "x-channel": "", + "x-api-key": credentials["aippt_access_key"], + "x-token": cls._get_api_token(credentials=credentials, user_id=user_id), + } + response = get(str(cls._api_base_url / "template_component" / "suit" / "select"), headers=headers) + + if response.status_code != 200: + raise Exception(f"Failed to connect to aippt: {response.text}") + + response = response.json() + + if response.get("code") != 0: + raise Exception(f"Failed to connect to aippt: {response.get('msg')}") + + colors = [ + { + "id": f"id-{item.get('id')}", + "name": item.get("name"), + "en_name": item.get("en_name", item.get("name")), + } + for item in response.get("data", {}).get("colour") or [] + ] + styles = [ + { + "id": f"id-{item.get('id')}", + "name": item.get("title"), + } + for item in response.get("data", {}).get("suit_style") or [] + ] + + with cls._style_cache_lock: + cls._style_cache[key] = {"colors": colors, "styles": styles, "expire": now + 60 * 60} + + return colors, styles + + def get_styles(self, user_id: str) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + """ + Get styles + + :param credentials: the credentials + :return: Tuple[list[dict[id, color]], list[dict[id, style]] + """ + if not self._tool.runtime.credentials.get("aippt_access_key") or not self._tool.runtime.credentials.get( + "aippt_secret_key" + ): + raise Exception("Please provide aippt credentials") + + return self._get_styles(credentials=self._tool.runtime.credentials, user_id=user_id) + + def _get_suit(self, style_id: int, colour_id: int) -> int: + """ + Get suit + """ + headers = { + "x-channel": "", + "x-api-key": self._tool.runtime.credentials["aippt_access_key"], + "x-token": self._get_api_token(credentials=self._tool.runtime.credentials, user_id="__dify_system__"), + } + response = get( + str(self._api_base_url / "template_component" / "suit" / "search"), + headers=headers, + params={"style_id": style_id, "colour_id": colour_id, "page": 1, "page_size": 1}, + ) + + if response.status_code != 200: + raise Exception(f"Failed to connect to aippt: {response.text}") + + response = response.json() + + if response.get("code") != 0: + raise Exception(f"Failed to connect to aippt: {response.get('msg')}") + + if len(response.get("data", {}).get("list") or []) > 0: + return response.get("data", {}).get("list")[0].get("id") + + raise Exception("Failed to get suit, the suit does not exist, please check the style and color") + + def get_runtime_parameters(self) -> list[ToolParameter]: + """ + Get runtime parameters + + Override this method to add runtime parameters to the tool. + """ + try: + colors, styles = self.get_styles(user_id="__dify_system__") + except Exception as e: + colors, styles = ( + [{"id": "-1", "name": "__default__", "en_name": "__default__"}], + [{"id": "-1", "name": "__default__", "en_name": "__default__"}], + ) + + return [ + ToolParameter( + name="color", + label=I18nObject(zh_Hans="颜色", en_US="Color"), + human_description=I18nObject(zh_Hans="颜色", en_US="Color"), + type=ToolParameter.ToolParameterType.SELECT, + form=ToolParameter.ToolParameterForm.FORM, + required=False, + default=colors[0]["id"], + options=[ + ToolParameterOption( + value=color["id"], label=I18nObject(zh_Hans=color["name"], en_US=color["en_name"]) + ) + for color in colors + ], + ), + ToolParameter( + name="style", + label=I18nObject(zh_Hans="风格", en_US="Style"), + human_description=I18nObject(zh_Hans="风格", en_US="Style"), + type=ToolParameter.ToolParameterType.SELECT, + form=ToolParameter.ToolParameterForm.FORM, + required=False, + default=styles[0]["id"], + options=[ + ToolParameterOption(value=style["id"], label=I18nObject(zh_Hans=style["name"], en_US=style["name"])) + for style in styles + ], + ), + ] + + +class AIPPTGenerateTool(BuiltinTool): + def __init__(self, **kwargs: Any): + super().__init__(**kwargs) + + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + return AIPPTGenerateToolAdapter(self)._invoke(user_id, tool_parameters) + + def get_runtime_parameters(self) -> list[ToolParameter]: + return AIPPTGenerateToolAdapter(self).get_runtime_parameters() + + @classmethod + def _get_api_token(cls, credentials: dict[str, str], user_id: str) -> str: + return AIPPTGenerateToolAdapter()._get_api_token(credentials, user_id) diff --git a/api/core/tools/provider/builtin/aippt/tools/aippt.yaml b/api/core/tools/provider/builtin/aippt/tools/aippt.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d35798ad66106e85f6ac078b4b18b56d3c4b917e --- /dev/null +++ b/api/core/tools/provider/builtin/aippt/tools/aippt.yaml @@ -0,0 +1,54 @@ +identity: + name: aippt + author: Dify + label: + en_US: AIPPT + zh_Hans: AIPPT +description: + human: + en_US: AI-generated PPT with one click, input your content topic, and let AI serve you one-stop + zh_Hans: AI一键生成PPT,输入你的内容主题,让AI为你一站式服务到底 + llm: A tool used to generate PPT with AI, input your content topic, and let AI generate PPT for you. +parameters: + - name: title + type: string + required: true + label: + en_US: Title + zh_Hans: 标题 + human_description: + en_US: The title of the PPT. + zh_Hans: PPT的标题。 + llm_description: The title of the PPT, which will be used to generate the PPT outline. + form: llm + - name: outline + type: string + required: false + label: + en_US: Outline + zh_Hans: 大纲 + human_description: + en_US: The outline of the PPT + zh_Hans: PPT的大纲 + llm_description: The outline of the PPT, which will be used to generate the PPT content. provide it if you have. + form: llm + - name: llm + type: select + required: true + label: + en_US: LLM model + zh_Hans: 生成大纲的LLM + options: + - value: aippt + label: + en_US: AIPPT default model + zh_Hans: AIPPT默认模型 + - value: wenxin + label: + en_US: Wenxin ErnieBot + zh_Hans: 文心一言 + default: aippt + human_description: + en_US: The LLM model used for generating PPT outline. + zh_Hans: 用于生成PPT大纲的LLM模型。 + form: form diff --git a/api/core/tools/provider/builtin/aliyuque/_assets/icon.svg b/api/core/tools/provider/builtin/aliyuque/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..82b23ebbc66e68ebcf363d77163b1ca85965b285 --- /dev/null +++ b/api/core/tools/provider/builtin/aliyuque/_assets/icon.svg @@ -0,0 +1,32 @@ + + 绿 lgo + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/aliyuque/aliyuque.py b/api/core/tools/provider/builtin/aliyuque/aliyuque.py new file mode 100644 index 0000000000000000000000000000000000000000..56eac1a4b570cfa7ccf6b299aac7da44926ba6f7 --- /dev/null +++ b/api/core/tools/provider/builtin/aliyuque/aliyuque.py @@ -0,0 +1,19 @@ +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.aliyuque.tools.base import AliYuqueTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class AliYuqueProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + token = credentials.get("token") + if not token: + raise ToolProviderCredentialValidationError("token is required") + + try: + resp = AliYuqueTool.auth(token) + if resp and resp.get("data", {}).get("id"): + return + + raise ToolProviderCredentialValidationError(resp) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/aliyuque/aliyuque.yaml b/api/core/tools/provider/builtin/aliyuque/aliyuque.yaml new file mode 100644 index 0000000000000000000000000000000000000000..73d39aa96cfd179f30056bfd2ff4b92fe723a18c --- /dev/null +++ b/api/core/tools/provider/builtin/aliyuque/aliyuque.yaml @@ -0,0 +1,29 @@ +identity: + author: 佐井 + name: aliyuque + label: + en_US: yuque + zh_Hans: 语雀 + pt_BR: yuque + description: + en_US: Yuque, https://www.yuque.com. + zh_Hans: 语雀,https://www.yuque.com。 + pt_BR: Yuque, https://www.yuque.com. + icon: icon.svg + tags: + - productivity + - search +credentials_for_provider: + token: + type: secret-input + required: true + label: + en_US: Yuque Team Token + zh_Hans: 语雀团队Token + placeholder: + en_US: Please input your Yuque team token + zh_Hans: 请输入你的语雀团队Token + help: + en_US: Get Alibaba Yuque team token + zh_Hans: 先获取语雀团队Token + url: https://www.yuque.com/settings/tokens diff --git a/api/core/tools/provider/builtin/aliyuque/tools/base.py b/api/core/tools/provider/builtin/aliyuque/tools/base.py new file mode 100644 index 0000000000000000000000000000000000000000..edfb9fea8ec4536d67e6c07aa8d23ab99cee99eb --- /dev/null +++ b/api/core/tools/provider/builtin/aliyuque/tools/base.py @@ -0,0 +1,42 @@ +from typing import Any + +import requests + + +class AliYuqueTool: + # yuque service url + server_url = "https://www.yuque.com" + + @staticmethod + def auth(token): + session = requests.Session() + session.headers.update({"Accept": "application/json", "X-Auth-Token": token}) + login = session.request("GET", AliYuqueTool.server_url + "/api/v2/user") + login.raise_for_status() + resp = login.json() + return resp + + def request(self, method: str, token, tool_parameters: dict[str, Any], path: str) -> str: + if not token: + raise Exception("token is required") + session = requests.Session() + session.headers.update({"accept": "application/json", "X-Auth-Token": token}) + new_params = {**tool_parameters} + + replacements = {k: v for k, v in new_params.items() if f"{{{k}}}" in path} + + for key, value in replacements.items(): + path = path.replace(f"{{{key}}}", str(value)) + del new_params[key] + + if method.upper() in {"POST", "PUT"}: + session.headers.update( + { + "Content-Type": "application/json", + } + ) + response = session.request(method.upper(), self.server_url + path, json=new_params) + else: + response = session.request(method, self.server_url + path, params=new_params) + response.raise_for_status() + return response.text diff --git a/api/core/tools/provider/builtin/aliyuque/tools/create_document.py b/api/core/tools/provider/builtin/aliyuque/tools/create_document.py new file mode 100644 index 0000000000000000000000000000000000000000..01080fd1d57f4d780b7689cae7ed9070af20d63a --- /dev/null +++ b/api/core/tools/provider/builtin/aliyuque/tools/create_document.py @@ -0,0 +1,15 @@ +from typing import Any, Union + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.aliyuque.tools.base import AliYuqueTool +from core.tools.tool.builtin_tool import BuiltinTool + + +class AliYuqueCreateDocumentTool(AliYuqueTool, BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + token = self.runtime.credentials.get("token", None) + if not token: + raise Exception("token is required") + return self.create_text_message(self.request("POST", token, tool_parameters, "/api/v2/repos/{book_id}/docs")) diff --git a/api/core/tools/provider/builtin/aliyuque/tools/create_document.yaml b/api/core/tools/provider/builtin/aliyuque/tools/create_document.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6ac8ae6696f33043b988ed7d351d5c004308a151 --- /dev/null +++ b/api/core/tools/provider/builtin/aliyuque/tools/create_document.yaml @@ -0,0 +1,99 @@ +identity: + name: aliyuque_create_document + author: 佐井 + label: + en_US: Create Document + zh_Hans: 创建文档 + icon: icon.svg +description: + human: + en_US: Creates a new document within a knowledge base without automatic addition to the table of contents. Requires a subsequent call to the "knowledge base directory update API". Supports setting visibility, format, and content. # 接口英文描述 + zh_Hans: 在知识库中创建新文档,但不会自动加入目录,需额外调用“知识库目录更新接口”。允许设置公开性、格式及正文内容。 + llm: Creates docs in a KB. + +parameters: + - name: book_id + type: string + required: true + form: llm + label: + en_US: Knowledge Base ID + zh_Hans: 知识库ID + human_description: + en_US: The unique identifier of the knowledge base where the document will be created. + zh_Hans: 文档将被创建的知识库的唯一标识。 + llm_description: ID of the target knowledge base. + + - name: title + type: string + required: false + form: llm + label: + en_US: Title + zh_Hans: 标题 + human_description: + en_US: The title of the document, defaults to 'Untitled' if not provided. + zh_Hans: 文档标题,默认为'无标题'如未提供。 + llm_description: Title of the document, defaults to 'Untitled'. + + - name: public + type: select + required: false + form: llm + options: + - value: 0 + label: + en_US: Private + zh_Hans: 私密 + - value: 1 + label: + en_US: Public + zh_Hans: 公开 + - value: 2 + label: + en_US: Enterprise-only + zh_Hans: 企业内公开 + label: + en_US: Visibility + zh_Hans: 公开性 + human_description: + en_US: Document visibility (0 Private, 1 Public, 2 Enterprise-only). + zh_Hans: 文档可见性(0 私密, 1 公开, 2 企业内公开)。 + llm_description: Doc visibility options, 0-private, 1-public, 2-enterprise. + + - name: format + type: select + required: false + form: llm + options: + - value: markdown + label: + en_US: markdown + zh_Hans: markdown + - value: html + label: + en_US: html + zh_Hans: html + - value: lake + label: + en_US: lake + zh_Hans: lake + label: + en_US: Content Format + zh_Hans: 内容格式 + human_description: + en_US: Format of the document content (markdown, HTML, Lake). + zh_Hans: 文档内容格式(markdown, HTML, Lake)。 + llm_description: Content format choices, markdown, HTML, Lake. + + - name: body + type: string + required: true + form: llm + label: + en_US: Body Content + zh_Hans: 正文内容 + human_description: + en_US: The actual content of the document. + zh_Hans: 文档的实际内容。 + llm_description: Content of the document. diff --git a/api/core/tools/provider/builtin/aliyuque/tools/delete_document.py b/api/core/tools/provider/builtin/aliyuque/tools/delete_document.py new file mode 100644 index 0000000000000000000000000000000000000000..84237cec30c56332a2256d124609e98ec9a4f3b9 --- /dev/null +++ b/api/core/tools/provider/builtin/aliyuque/tools/delete_document.py @@ -0,0 +1,17 @@ +from typing import Any, Union + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.aliyuque.tools.base import AliYuqueTool +from core.tools.tool.builtin_tool import BuiltinTool + + +class AliYuqueDeleteDocumentTool(AliYuqueTool, BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + token = self.runtime.credentials.get("token", None) + if not token: + raise Exception("token is required") + return self.create_text_message( + self.request("DELETE", token, tool_parameters, "/api/v2/repos/{book_id}/docs/{id}") + ) diff --git a/api/core/tools/provider/builtin/aliyuque/tools/delete_document.yaml b/api/core/tools/provider/builtin/aliyuque/tools/delete_document.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dddd62d3048c350db16a359199d2c7034317ae82 --- /dev/null +++ b/api/core/tools/provider/builtin/aliyuque/tools/delete_document.yaml @@ -0,0 +1,37 @@ +identity: + name: aliyuque_delete_document + author: 佐井 + label: + en_US: Delete Document + zh_Hans: 删除文档 + icon: icon.svg +description: + human: + en_US: Delete Document + zh_Hans: 根据id删除文档 + llm: Delete document. + +parameters: + - name: book_id + type: string + required: true + form: llm + label: + en_US: Knowledge Base ID + zh_Hans: 知识库ID + human_description: + en_US: The unique identifier of the knowledge base where the document will be created. + zh_Hans: 文档将被创建的知识库的唯一标识。 + llm_description: ID of the target knowledge base. + + - name: id + type: string + required: true + form: llm + label: + en_US: Document ID or Path + zh_Hans: 文档 ID or 路径 + human_description: + en_US: Document ID or path. + zh_Hans: 文档 ID or 路径。 + llm_description: Document ID or path. diff --git a/api/core/tools/provider/builtin/aliyuque/tools/describe_book_index_page.py b/api/core/tools/provider/builtin/aliyuque/tools/describe_book_index_page.py new file mode 100644 index 0000000000000000000000000000000000000000..c23d30059a8424d7d23a70a2ee49b3f8c68a0a2b --- /dev/null +++ b/api/core/tools/provider/builtin/aliyuque/tools/describe_book_index_page.py @@ -0,0 +1,17 @@ +from typing import Any, Union + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.aliyuque.tools.base import AliYuqueTool +from core.tools.tool.builtin_tool import BuiltinTool + + +class AliYuqueDescribeBookIndexPageTool(AliYuqueTool, BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + token = self.runtime.credentials.get("token", None) + if not token: + raise Exception("token is required") + return self.create_text_message( + self.request("GET", token, tool_parameters, "/api/v2/repos/{group_login}/{book_slug}/index_page") + ) diff --git a/api/core/tools/provider/builtin/aliyuque/tools/describe_book_index_page.yaml b/api/core/tools/provider/builtin/aliyuque/tools/describe_book_index_page.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5e490725d1888269fa31c79736059d884914208b --- /dev/null +++ b/api/core/tools/provider/builtin/aliyuque/tools/describe_book_index_page.yaml @@ -0,0 +1,38 @@ +identity: + name: aliyuque_describe_book_index_page + author: 佐井 + label: + en_US: Get Repo Index Page + zh_Hans: 获取知识库首页 + icon: icon.svg + +description: + human: + en_US: Retrieves the homepage of a knowledge base within a group, supporting both book ID and group login with book slug access. + zh_Hans: 获取团队中知识库的首页信息,可通过书籍ID或团队登录名与书籍路径访问。 + llm: Fetches the knowledge base homepage using group and book identifiers with support for alternate access paths. + +parameters: + - name: group_login + type: string + required: true + form: llm + label: + en_US: Group Login + zh_Hans: 团队登录名 + human_description: + en_US: The login name of the group that owns the knowledge base. + zh_Hans: 拥有该知识库的团队登录名。 + llm_description: Team login identifier for the knowledge base owner. + + - name: book_slug + type: string + required: true + form: llm + label: + en_US: Book Slug + zh_Hans: 知识库路径 + human_description: + en_US: The unique slug representing the path of the knowledge base. + zh_Hans: 知识库的唯一路径标识。 + llm_description: Unique path identifier for the knowledge base. diff --git a/api/core/tools/provider/builtin/aliyuque/tools/describe_book_table_of_contents.py b/api/core/tools/provider/builtin/aliyuque/tools/describe_book_table_of_contents.py new file mode 100644 index 0000000000000000000000000000000000000000..36f8c10d6fd79d7d9caa98d19466b174e6c34550 --- /dev/null +++ b/api/core/tools/provider/builtin/aliyuque/tools/describe_book_table_of_contents.py @@ -0,0 +1,15 @@ +from typing import Any, Union + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.aliyuque.tools.base import AliYuqueTool +from core.tools.tool.builtin_tool import BuiltinTool + + +class YuqueDescribeBookTableOfContentsTool(AliYuqueTool, BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> (Union)[ToolInvokeMessage, list[ToolInvokeMessage]]: + token = self.runtime.credentials.get("token", None) + if not token: + raise Exception("token is required") + return self.create_text_message(self.request("GET", token, tool_parameters, "/api/v2/repos/{book_id}/toc")) diff --git a/api/core/tools/provider/builtin/aliyuque/tools/describe_book_table_of_contents.yaml b/api/core/tools/provider/builtin/aliyuque/tools/describe_book_table_of_contents.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0a481b59ebedad1b5eba19289b28f65a6d0729f3 --- /dev/null +++ b/api/core/tools/provider/builtin/aliyuque/tools/describe_book_table_of_contents.yaml @@ -0,0 +1,25 @@ +identity: + name: aliyuque_describe_book_table_of_contents + author: 佐井 + label: + en_US: Get Book's Table of Contents + zh_Hans: 获取知识库的目录 + icon: icon.svg +description: + human: + en_US: Get Book's Table of Contents. + zh_Hans: 获取知识库的目录。 + llm: Get Book's Table of Contents. + +parameters: + - name: book_id + type: string + required: true + form: llm + label: + en_US: Book ID + zh_Hans: 知识库 ID + human_description: + en_US: Book ID. + zh_Hans: 知识库 ID。 + llm_description: Book ID. diff --git a/api/core/tools/provider/builtin/aliyuque/tools/describe_document_content.py b/api/core/tools/provider/builtin/aliyuque/tools/describe_document_content.py new file mode 100644 index 0000000000000000000000000000000000000000..a69bf121f7e5aee4d93f23e6459f2fbfff5ba02d --- /dev/null +++ b/api/core/tools/provider/builtin/aliyuque/tools/describe_document_content.py @@ -0,0 +1,53 @@ +import json +from typing import Any, Union +from urllib.parse import urlparse + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.aliyuque.tools.base import AliYuqueTool +from core.tools.tool.builtin_tool import BuiltinTool + + +class AliYuqueDescribeDocumentContentTool(AliYuqueTool, BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + new_params = {**tool_parameters} + token = new_params.pop("token") + if not token or token.lower() == "none": + token = self.runtime.credentials.get("token", None) + if not token: + raise Exception("token is required") + new_params = {**tool_parameters} + url = new_params.pop("url") + if not url or not url.startswith("http"): + raise Exception("url is not valid") + + parsed_url = urlparse(url) + path_parts = parsed_url.path.strip("/").split("/") + if len(path_parts) < 3: + raise Exception("url is not correct") + doc_id = path_parts[-1] + book_slug = path_parts[-2] + group_id = path_parts[-3] + + new_params["group_login"] = group_id + new_params["book_slug"] = book_slug + index_page = json.loads( + self.request("GET", token, new_params, "/api/v2/repos/{group_login}/{book_slug}/index_page") + ) + book_id = index_page.get("data", {}).get("book", {}).get("id") + if not book_id: + raise Exception(f"can not parse book_id from {index_page}") + + new_params["book_id"] = book_id + new_params["id"] = doc_id + data = self.request("GET", token, new_params, "/api/v2/repos/{book_id}/docs/{id}") + data = json.loads(data) + body_only = tool_parameters.get("body_only") or "" + if body_only.lower() == "true": + return self.create_text_message(data.get("data").get("body")) + else: + raw = data.get("data") + del raw["body_lake"] + del raw["body_html"] + return self.create_text_message(json.dumps(data)) diff --git a/api/core/tools/provider/builtin/aliyuque/tools/describe_document_content.yaml b/api/core/tools/provider/builtin/aliyuque/tools/describe_document_content.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6116886a96b790dfb40ea2dad340d7f179b589ee --- /dev/null +++ b/api/core/tools/provider/builtin/aliyuque/tools/describe_document_content.yaml @@ -0,0 +1,50 @@ +identity: + name: aliyuque_describe_document_content + author: 佐井 + label: + en_US: Fetch Document Content + zh_Hans: 获取文档内容 + icon: icon.svg + +description: + human: + en_US: Retrieves document content from Yuque based on the provided document URL, which can be a normal or shared link. + zh_Hans: 根据提供的语雀文档地址(支持正常链接或分享链接)获取文档内容。 + llm: Fetches Yuque document content given a URL. + +parameters: + - name: url + type: string + required: true + form: llm + label: + en_US: Document URL + zh_Hans: 文档地址 + human_description: + en_US: The URL of the document to retrieve content from, can be normal or shared. + zh_Hans: 需要获取内容的文档地址,可以是正常链接或分享链接。 + llm_description: URL of the Yuque document to fetch content. + + - name: body_only + type: string + required: false + form: llm + label: + en_US: return body content only + zh_Hans: 仅返回body内容 + human_description: + en_US: true:Body content only, false:Full response with metadata. + zh_Hans: true:仅返回body内容,不返回其他元数据,false:返回所有元数据。 + llm_description: true:Body content only, false:Full response with metadata. + + - name: token + type: secret-input + required: false + form: llm + label: + en_US: Yuque API Token + zh_Hans: 语雀接口Token + human_description: + en_US: The token for calling the Yuque API defaults to the Yuque token bound to the current tool if not provided. + zh_Hans: 调用语雀接口的token,如果不传则默认为当前工具绑定的语雀Token。 + llm_description: If the token for calling the Yuque API is not provided, it will default to the Yuque token bound to the current tool. diff --git a/api/core/tools/provider/builtin/aliyuque/tools/describe_documents.py b/api/core/tools/provider/builtin/aliyuque/tools/describe_documents.py new file mode 100644 index 0000000000000000000000000000000000000000..7a45684bed04984a94b1d79ddee46b9eabe52682 --- /dev/null +++ b/api/core/tools/provider/builtin/aliyuque/tools/describe_documents.py @@ -0,0 +1,17 @@ +from typing import Any, Union + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.aliyuque.tools.base import AliYuqueTool +from core.tools.tool.builtin_tool import BuiltinTool + + +class AliYuqueDescribeDocumentsTool(AliYuqueTool, BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + token = self.runtime.credentials.get("token", None) + if not token: + raise Exception("token is required") + return self.create_text_message( + self.request("GET", token, tool_parameters, "/api/v2/repos/{book_id}/docs/{id}") + ) diff --git a/api/core/tools/provider/builtin/aliyuque/tools/describe_documents.yaml b/api/core/tools/provider/builtin/aliyuque/tools/describe_documents.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0b14c1afba684dcedce4fbfc65fc5084f9e401da --- /dev/null +++ b/api/core/tools/provider/builtin/aliyuque/tools/describe_documents.yaml @@ -0,0 +1,38 @@ +identity: + name: aliyuque_describe_documents + author: 佐井 + label: + en_US: Get Doc Detail + zh_Hans: 获取文档详情 + icon: icon.svg + +description: + human: + en_US: Retrieves detailed information of a specific document identified by its ID or path within a knowledge base. + zh_Hans: 根据知识库ID和文档ID或路径获取文档详细信息。 + llm: Fetches detailed doc info using ID/path from a knowledge base; supports doc lookup in Yuque. + +parameters: + - name: book_id + type: string + required: true + form: llm + label: + en_US: Knowledge Base ID + zh_Hans: 知识库 ID + human_description: + en_US: Identifier for the knowledge base where the document resides. + zh_Hans: 文档所属知识库的唯一标识。 + llm_description: ID of the knowledge base holding the document. + + - name: id + type: string + required: true + form: llm + label: + en_US: Document ID or Path + zh_Hans: 文档 ID 或路径 + human_description: + en_US: The unique identifier or path of the document to retrieve. + zh_Hans: 需要获取的文档的ID或其在知识库中的路径。 + llm_description: Unique doc ID or its path for retrieval. diff --git a/api/core/tools/provider/builtin/aliyuque/tools/update_book_table_of_contents.py b/api/core/tools/provider/builtin/aliyuque/tools/update_book_table_of_contents.py new file mode 100644 index 0000000000000000000000000000000000000000..ca0a3909f807094a66c0b38b38d433ae36473599 --- /dev/null +++ b/api/core/tools/provider/builtin/aliyuque/tools/update_book_table_of_contents.py @@ -0,0 +1,21 @@ +from typing import Any, Union + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.aliyuque.tools.base import AliYuqueTool +from core.tools.tool.builtin_tool import BuiltinTool + + +class YuqueDescribeBookTableOfContentsTool(AliYuqueTool, BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> (Union)[ToolInvokeMessage, list[ToolInvokeMessage]]: + token = self.runtime.credentials.get("token", None) + if not token: + raise Exception("token is required") + + doc_ids = tool_parameters.get("doc_ids") + if doc_ids: + doc_ids = [int(doc_id.strip()) for doc_id in doc_ids.split(",")] + tool_parameters["doc_ids"] = doc_ids + + return self.create_text_message(self.request("PUT", token, tool_parameters, "/api/v2/repos/{book_id}/toc")) diff --git a/api/core/tools/provider/builtin/aliyuque/tools/update_book_table_of_contents.yaml b/api/core/tools/provider/builtin/aliyuque/tools/update_book_table_of_contents.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f85970348b1f13e47b8a69d1ad298f4b5fb0615d --- /dev/null +++ b/api/core/tools/provider/builtin/aliyuque/tools/update_book_table_of_contents.yaml @@ -0,0 +1,222 @@ +identity: + name: aliyuque_update_book_table_of_contents + author: 佐井 + label: + en_US: Update Book's Table of Contents + zh_Hans: 更新知识库目录 + icon: icon.svg +description: + human: + en_US: Update Book's Table of Contents. + zh_Hans: 更新知识库目录。 + llm: Update Book's Table of Contents. + +parameters: + - name: book_id + type: string + required: true + form: llm + label: + en_US: Book ID + zh_Hans: 知识库 ID + human_description: + en_US: Book ID. + zh_Hans: 知识库 ID。 + llm_description: Book ID. + + - name: action + type: select + required: true + form: llm + options: + - value: appendNode + label: + en_US: appendNode + zh_Hans: appendNode + pt_BR: appendNode + - value: prependNode + label: + en_US: prependNode + zh_Hans: prependNode + pt_BR: prependNode + - value: editNode + label: + en_US: editNode + zh_Hans: editNode + pt_BR: editNode + - value: editNode + label: + en_US: removeNode + zh_Hans: removeNode + pt_BR: removeNode + label: + en_US: Action Type + zh_Hans: 操作 + human_description: + en_US: In the operation scenario, sibling node prepending is not supported, deleting a node doesn't remove associated documents, and node deletion has two modes, 'sibling' (delete current node) and 'child' (delete current node and its children). + zh_Hans: 操作,创建场景下不支持同级头插 prependNode,删除节点不会删除关联文档,删除节点时action_mode=sibling (删除当前节点), action_mode=child (删除当前节点及子节点) + llm_description: In the operation scenario, sibling node prepending is not supported, deleting a node doesn't remove associated documents, and node deletion has two modes, 'sibling' (delete current node) and 'child' (delete current node and its children). + + + - name: action_mode + type: select + required: false + form: llm + options: + - value: sibling + label: + en_US: sibling + zh_Hans: 同级 + pt_BR: sibling + - value: child + label: + en_US: child + zh_Hans: 子集 + pt_BR: child + label: + en_US: Action Type + zh_Hans: 操作 + human_description: + en_US: Operation mode (sibling:same level, child:child level). + zh_Hans: 操作模式 (sibling:同级, child:子级)。 + llm_description: Operation mode (sibling:same level, child:child level). + + - name: target_uuid + type: string + required: false + form: llm + label: + en_US: Target node UUID + zh_Hans: 目标节点 UUID + human_description: + en_US: Target node UUID, defaults to root node if left empty. + zh_Hans: 目标节点 UUID, 不填默认为根节点。 + llm_description: Target node UUID, defaults to root node if left empty. + + - name: node_uuid + type: string + required: false + form: llm + label: + en_US: Node UUID + zh_Hans: 操作节点 UUID + human_description: + en_US: Operation node UUID [required for move/update/delete]. + zh_Hans: 操作节点 UUID [移动/更新/删除必填]。 + llm_description: Operation node UUID [required for move/update/delete]. + + - name: doc_ids + type: string + required: false + form: llm + label: + en_US: Document IDs + zh_Hans: 文档id列表 + human_description: + en_US: Document IDs [required for creating documents], separate multiple IDs with ','. + zh_Hans: 文档 IDs [创建文档必填],多个用','分隔。 + llm_description: Document IDs [required for creating documents], separate multiple IDs with ','. + + + - name: type + type: select + required: false + form: llm + default: DOC + options: + - value: DOC + label: + en_US: DOC + zh_Hans: 文档 + pt_BR: DOC + - value: LINK + label: + en_US: LINK + zh_Hans: 链接 + pt_BR: LINK + - value: TITLE + label: + en_US: TITLE + zh_Hans: 分组 + pt_BR: TITLE + label: + en_US: Node type + zh_Hans: 操节点类型 + human_description: + en_US: Node type [required for creation] (DOC:document, LINK:external link, TITLE:group). + zh_Hans: 操节点类型 [创建必填] (DOC:文档, LINK:外链, TITLE:分组)。 + llm_description: Node type [required for creation] (DOC:document, LINK:external link, TITLE:group). + + - name: title + type: string + required: false + form: llm + label: + en_US: Node Name + zh_Hans: 节点名称 + human_description: + en_US: Node name [required for creating groups/external links]. + zh_Hans: 节点名称 [创建分组/外链必填]。 + llm_description: Node name [required for creating groups/external links]. + + - name: url + type: string + required: false + form: llm + label: + en_US: Node URL + zh_Hans: 节点URL + human_description: + en_US: Node URL [required for creating external links]. + zh_Hans: 节点 URL [创建外链必填]。 + llm_description: Node URL [required for creating external links]. + + + - name: open_window + type: select + required: false + form: llm + default: 0 + options: + - value: 0 + label: + en_US: DOC + zh_Hans: Current Page + pt_BR: DOC + - value: 1 + label: + en_US: LINK + zh_Hans: New Page + pt_BR: LINK + label: + en_US: Open in new window + zh_Hans: 是否新窗口打开 + human_description: + en_US: Open in new window [optional for external links] (0:open in current page, 1:open in new window). + zh_Hans: 是否新窗口打开 [外链选填] (0:当前页打开, 1:新窗口打开)。 + llm_description: Open in new window [optional for external links] (0:open in current page, 1:open in new window). + + + - name: visible + type: select + required: false + form: llm + default: 1 + options: + - value: 0 + label: + en_US: Invisible + zh_Hans: 隐藏 + pt_BR: Invisible + - value: 1 + label: + en_US: Visible + zh_Hans: 可见 + pt_BR: Visible + label: + en_US: Visibility + zh_Hans: 是否可见 + human_description: + en_US: Visibility (0:invisible, 1:visible). + zh_Hans: 是否可见 (0:不可见, 1:可见)。 + llm_description: Visibility (0:invisible, 1:visible). diff --git a/api/core/tools/provider/builtin/aliyuque/tools/update_document.py b/api/core/tools/provider/builtin/aliyuque/tools/update_document.py new file mode 100644 index 0000000000000000000000000000000000000000..d7eba46ad968dd8f07c6adb035780e43e2e2731c --- /dev/null +++ b/api/core/tools/provider/builtin/aliyuque/tools/update_document.py @@ -0,0 +1,17 @@ +from typing import Any, Union + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.aliyuque.tools.base import AliYuqueTool +from core.tools.tool.builtin_tool import BuiltinTool + + +class AliYuqueUpdateDocumentTool(AliYuqueTool, BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + token = self.runtime.credentials.get("token", None) + if not token: + raise Exception("token is required") + return self.create_text_message( + self.request("PUT", token, tool_parameters, "/api/v2/repos/{book_id}/docs/{id}") + ) diff --git a/api/core/tools/provider/builtin/aliyuque/tools/update_document.yaml b/api/core/tools/provider/builtin/aliyuque/tools/update_document.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c2da6b179acdd98a920fddfd9ce263d1ae535407 --- /dev/null +++ b/api/core/tools/provider/builtin/aliyuque/tools/update_document.yaml @@ -0,0 +1,87 @@ +identity: + name: aliyuque_update_document + author: 佐井 + label: + en_US: Update Document + zh_Hans: 更新文档 + icon: icon.svg +description: + human: + en_US: Update an existing document within a specified knowledge base by providing the document ID or path. + zh_Hans: 通过提供文档ID或路径,更新指定知识库中的现有文档。 + llm: Update doc in a knowledge base via ID/path. +parameters: + - name: book_id + type: string + required: true + form: llm + label: + en_US: Knowledge Base ID + zh_Hans: 知识库 ID + human_description: + en_US: The unique identifier of the knowledge base where the document resides. + zh_Hans: 文档所属知识库的ID。 + llm_description: ID of the knowledge base holding the doc. + - name: id + type: string + required: true + form: llm + label: + en_US: Document ID or Path + zh_Hans: 文档 ID 或 路径 + human_description: + en_US: The unique identifier or the path of the document to be updated. + zh_Hans: 要更新的文档的唯一ID或路径。 + llm_description: Doc's ID or path for update. + + - name: title + type: string + required: false + form: llm + label: + en_US: Title + zh_Hans: 标题 + human_description: + en_US: The title of the document, defaults to 'Untitled' if not provided. + zh_Hans: 文档标题,默认为'无标题'如未提供。 + llm_description: Title of the document, defaults to 'Untitled'. + + - name: format + type: select + required: false + form: llm + options: + - value: markdown + label: + en_US: markdown + zh_Hans: markdown + pt_BR: markdown + - value: html + label: + en_US: html + zh_Hans: html + pt_BR: html + - value: lake + label: + en_US: lake + zh_Hans: lake + pt_BR: lake + label: + en_US: Content Format + zh_Hans: 内容格式 + human_description: + en_US: Format of the document content (markdown, HTML, Lake). + zh_Hans: 文档内容格式(markdown, HTML, Lake)。 + llm_description: Content format choices, markdown, HTML, Lake. + + - name: body + type: string + required: true + form: llm + label: + en_US: Body Content + zh_Hans: 正文内容 + human_description: + en_US: The actual content of the document. + zh_Hans: 文档的实际内容。 + llm_description: Content of the document. diff --git a/api/core/tools/provider/builtin/alphavantage/_assets/icon.svg b/api/core/tools/provider/builtin/alphavantage/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..785432943bc14819b58390d0c5e91de12729a21b --- /dev/null +++ b/api/core/tools/provider/builtin/alphavantage/_assets/icon.svg @@ -0,0 +1,7 @@ + + + 形状结合 + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/alphavantage/alphavantage.py b/api/core/tools/provider/builtin/alphavantage/alphavantage.py new file mode 100644 index 0000000000000000000000000000000000000000..a84630e5aa990abbba59c9a4851478d07b09da11 --- /dev/null +++ b/api/core/tools/provider/builtin/alphavantage/alphavantage.py @@ -0,0 +1,22 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.alphavantage.tools.query_stock import QueryStockTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class AlphaVantageProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + QueryStockTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "code": "AAPL", # Apple Inc. + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/alphavantage/alphavantage.yaml b/api/core/tools/provider/builtin/alphavantage/alphavantage.yaml new file mode 100644 index 0000000000000000000000000000000000000000..710510cfd8ed4af8869c7655fb9368367faf13f1 --- /dev/null +++ b/api/core/tools/provider/builtin/alphavantage/alphavantage.yaml @@ -0,0 +1,31 @@ +identity: + author: zhuhao + name: alphavantage + label: + en_US: AlphaVantage + zh_Hans: AlphaVantage + pt_BR: AlphaVantage + description: + en_US: AlphaVantage is an online platform that provides financial market data and APIs, making it convenient for individual investors and developers to access stock quotes, technical indicators, and stock analysis. + zh_Hans: AlphaVantage是一个在线平台,它提供金融市场数据和API,便于个人投资者和开发者获取股票报价、技术指标和股票分析。 + pt_BR: AlphaVantage is an online platform that provides financial market data and APIs, making it convenient for individual investors and developers to access stock quotes, technical indicators, and stock analysis. + icon: icon.svg + tags: + - finance +credentials_for_provider: + api_key: + type: secret-input + required: true + label: + en_US: AlphaVantage API key + zh_Hans: AlphaVantage API key + pt_BR: AlphaVantage API key + placeholder: + en_US: Please input your AlphaVantage API key + zh_Hans: 请输入你的 AlphaVantage API key + pt_BR: Please input your AlphaVantage API key + help: + en_US: Get your AlphaVantage API key from AlphaVantage + zh_Hans: 从 AlphaVantage 获取您的 AlphaVantage API key + pt_BR: Get your AlphaVantage API key from AlphaVantage + url: https://www.alphavantage.co/support/#api-key diff --git a/api/core/tools/provider/builtin/alphavantage/tools/query_stock.py b/api/core/tools/provider/builtin/alphavantage/tools/query_stock.py new file mode 100644 index 0000000000000000000000000000000000000000..d06611acd05d1d339e063eaf26be120bc45ec86f --- /dev/null +++ b/api/core/tools/provider/builtin/alphavantage/tools/query_stock.py @@ -0,0 +1,48 @@ +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +ALPHAVANTAGE_API_URL = "https://www.alphavantage.co/query" + + +class QueryStockTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + stock_code = tool_parameters.get("code", "") + if not stock_code: + return self.create_text_message("Please tell me your stock code") + + if "api_key" not in self.runtime.credentials or not self.runtime.credentials.get("api_key"): + return self.create_text_message("Alpha Vantage API key is required.") + + params = { + "function": "TIME_SERIES_DAILY", + "symbol": stock_code, + "outputsize": "compact", + "datatype": "json", + "apikey": self.runtime.credentials["api_key"], + } + response = requests.get(url=ALPHAVANTAGE_API_URL, params=params) + response.raise_for_status() + result = self._handle_response(response.json()) + return self.create_json_message(result) + + def _handle_response(self, response: dict[str, Any]) -> dict[str, Any]: + result = response.get("Time Series (Daily)", {}) + if not result: + return {} + stock_result = {} + for k, v in result.items(): + stock_result[k] = {} + stock_result[k]["open"] = v.get("1. open") + stock_result[k]["high"] = v.get("2. high") + stock_result[k]["low"] = v.get("3. low") + stock_result[k]["close"] = v.get("4. close") + stock_result[k]["volume"] = v.get("5. volume") + return stock_result diff --git a/api/core/tools/provider/builtin/alphavantage/tools/query_stock.yaml b/api/core/tools/provider/builtin/alphavantage/tools/query_stock.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d89f34e373f9fa6045cced012ecf89bd36725eaa --- /dev/null +++ b/api/core/tools/provider/builtin/alphavantage/tools/query_stock.yaml @@ -0,0 +1,27 @@ +identity: + name: query_stock + author: zhuhao + label: + en_US: query_stock + zh_Hans: query_stock + pt_BR: query_stock +description: + human: + en_US: Retrieve information such as daily opening price, daily highest price, daily lowest price, daily closing price, and daily trading volume for a specified stock symbol. + zh_Hans: 获取指定股票代码的每日开盘价、每日最高价、每日最低价、每日收盘价和每日交易量等信息。 + pt_BR: Retrieve information such as daily opening price, daily highest price, daily lowest price, daily closing price, and daily trading volume for a specified stock symbol + llm: Retrieve information such as daily opening price, daily highest price, daily lowest price, daily closing price, and daily trading volume for a specified stock symbol +parameters: + - name: code + type: string + required: true + label: + en_US: stock code + zh_Hans: 股票代码 + pt_BR: stock code + human_description: + en_US: stock code + zh_Hans: 股票代码 + pt_BR: stock code + llm_description: stock code for query from alphavantage + form: llm diff --git a/api/core/tools/provider/builtin/arxiv/_assets/icon.svg b/api/core/tools/provider/builtin/arxiv/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..0e60f635739993fa0a64bbd2fdaf8458ad87780a --- /dev/null +++ b/api/core/tools/provider/builtin/arxiv/_assets/icon.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/arxiv/arxiv.py b/api/core/tools/provider/builtin/arxiv/arxiv.py new file mode 100644 index 0000000000000000000000000000000000000000..ebb2d1a8c47be93f0c84625ab74998c2beba5e5c --- /dev/null +++ b/api/core/tools/provider/builtin/arxiv/arxiv.py @@ -0,0 +1,20 @@ +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.arxiv.tools.arxiv_search import ArxivSearchTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class ArxivProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + try: + ArxivSearchTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "query": "John Doe", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/arxiv/arxiv.yaml b/api/core/tools/provider/builtin/arxiv/arxiv.yaml new file mode 100644 index 0000000000000000000000000000000000000000..25aec97bb795e3e9fe5c25bc0614742d795ba5be --- /dev/null +++ b/api/core/tools/provider/builtin/arxiv/arxiv.yaml @@ -0,0 +1,14 @@ +identity: + author: Yash Parmar + name: arxiv + label: + en_US: ArXiv + zh_Hans: ArXiv + ja_JP: ArXiv + description: + en_US: Access to a vast repository of scientific papers and articles in various fields of research. + zh_Hans: 访问各个研究领域大量科学论文和文章的存储库。 + ja_JP: 多様な研究分野の科学論文や記事の膨大なリポジトリへのアクセス。 + icon: icon.svg + tags: + - search diff --git a/api/core/tools/provider/builtin/arxiv/tools/arxiv_search.py b/api/core/tools/provider/builtin/arxiv/tools/arxiv_search.py new file mode 100644 index 0000000000000000000000000000000000000000..8bd16050ecf0a67fef61836d7b97f1f021cddfae --- /dev/null +++ b/api/core/tools/provider/builtin/arxiv/tools/arxiv_search.py @@ -0,0 +1,119 @@ +import logging +from typing import Any, Optional + +import arxiv # type: ignore +from pydantic import BaseModel, Field + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +logger = logging.getLogger(__name__) + + +class ArxivAPIWrapper(BaseModel): + """Wrapper around ArxivAPI. + + To use, you should have the ``arxiv`` python package installed. + https://lukasschwab.me/arxiv.py/index.html + This wrapper will use the Arxiv API to conduct searches and + fetch document summaries. By default, it will return the document summaries + of the top-k results. + It limits the Document content by doc_content_chars_max. + Set doc_content_chars_max=None if you don't want to limit the content size. + + Args: + top_k_results: number of the top-scored document used for the arxiv tool + ARXIV_MAX_QUERY_LENGTH: the cut limit on the query used for the arxiv tool. + load_max_docs: a limit to the number of loaded documents + load_all_available_meta: + if True: the `metadata` of the loaded Documents contains all available + meta info (see https://lukasschwab.me/arxiv.py/index.html#Result), + if False: the `metadata` contains only the published date, title, + authors and summary. + doc_content_chars_max: an optional cut limit for the length of a document's + content + + Example: + .. code-block:: python + + arxiv = ArxivAPIWrapper( + top_k_results = 3, + ARXIV_MAX_QUERY_LENGTH = 300, + load_max_docs = 3, + load_all_available_meta = False, + doc_content_chars_max = 40000 + ) + arxiv.run("tree of thought llm) + """ + + arxiv_search: type[arxiv.Search] = arxiv.Search #: :meta private: + arxiv_http_error: tuple[type[Exception]] = (arxiv.ArxivError, arxiv.UnexpectedEmptyPageError, arxiv.HTTPError) + top_k_results: int = 3 + ARXIV_MAX_QUERY_LENGTH: int = 300 + load_max_docs: int = 100 + load_all_available_meta: bool = False + doc_content_chars_max: Optional[int] = 4000 + + def run(self, query: str) -> str: + """ + Performs an arxiv search and A single string + with the publish date, title, authors, and summary + for each article separated by two newlines. + + If an error occurs or no documents found, error text + is returned instead. Wrapper for + https://lukasschwab.me/arxiv.py/index.html#Search + + Args: + query: a plaintext search query + """ + try: + results = self.arxiv_search( # type: ignore + query[: self.ARXIV_MAX_QUERY_LENGTH], max_results=self.top_k_results + ).results() + except arxiv_http_error as ex: + return f"Arxiv exception: {ex}" + docs = [ + f"Published: {result.updated.date()}\n" + f"Title: {result.title}\n" + f"Authors: {', '.join(a.name for a in result.authors)}\n" + f"Summary: {result.summary}" + for result in results + ] + if docs: + return "\n\n".join(docs)[: self.doc_content_chars_max] + else: + return "No good Arxiv Result was found" + + +class ArxivSearchInput(BaseModel): + query: str = Field(..., description="Search query.") + + +class ArxivSearchTool(BuiltinTool): + """ + A tool for searching articles on Arxiv. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage | list[ToolInvokeMessage]: + """ + Invokes the Arxiv search tool with the given user ID and tool parameters. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (dict[str, Any]): The parameters for the tool, including the 'query' parameter. + + Returns: + ToolInvokeMessage | list[ToolInvokeMessage]: The result of the tool invocation, + which can be a single message or a list of messages. + """ + query = tool_parameters.get("query", "") + + if not query: + return self.create_text_message("Please input query") + + arxiv = ArxivAPIWrapper() + + response = arxiv.run(query) + + return self.create_text_message(self.summary(user_id=user_id, content=response)) diff --git a/api/core/tools/provider/builtin/arxiv/tools/arxiv_search.yaml b/api/core/tools/provider/builtin/arxiv/tools/arxiv_search.yaml new file mode 100644 index 0000000000000000000000000000000000000000..afc1925df3b45ab6dd6252e8ef01b4cc9d7e1767 --- /dev/null +++ b/api/core/tools/provider/builtin/arxiv/tools/arxiv_search.yaml @@ -0,0 +1,27 @@ +identity: + name: arxiv_search + author: Yash Parmar + label: + en_US: Arxiv Search + zh_Hans: Arxiv 搜索 + ja_JP: Arxiv 検索 +description: + human: + en_US: A tool for searching scientific papers and articles from the Arxiv repository. Input can be an Arxiv ID or an author's name. + zh_Hans: 一个用于从Arxiv存储库搜索科学论文和文章的工具。 输入可以是Arxiv ID或作者姓名。 + ja_JP: Arxivリポジトリから科学論文や記事を検索するためのツールです。入力はArxiv IDまたは著者名にすることができます。 + llm: A tool for searching scientific papers and articles from the Arxiv repository. Input can be an Arxiv ID or an author's name. +parameters: + - name: query + type: string + required: true + label: + en_US: Query string + zh_Hans: 查询字符串 + ja_JP: クエリ文字列 + human_description: + en_US: The Arxiv ID or author's name used for searching. + zh_Hans: 用于搜索的Arxiv ID或作者姓名。 + ja_JP: 検索に使用されるArxiv IDまたは著者名。 + llm_description: The Arxiv ID or author's name used for searching. + form: llm diff --git a/api/core/tools/provider/builtin/audio/_assets/icon.svg b/api/core/tools/provider/builtin/audio/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..08cc4ede66b90ea6de6af63eab9fcf5a27666d5f --- /dev/null +++ b/api/core/tools/provider/builtin/audio/_assets/icon.svg @@ -0,0 +1,3 @@ + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/audio/audio.py b/api/core/tools/provider/builtin/audio/audio.py new file mode 100644 index 0000000000000000000000000000000000000000..1f15386f78a0580c7e3e672d24c4de2efd08beb8 --- /dev/null +++ b/api/core/tools/provider/builtin/audio/audio.py @@ -0,0 +1,6 @@ +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class AudioToolProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + pass diff --git a/api/core/tools/provider/builtin/audio/audio.yaml b/api/core/tools/provider/builtin/audio/audio.yaml new file mode 100644 index 0000000000000000000000000000000000000000..07db268dacc446b8b66b8ac76e1888a9d485f3df --- /dev/null +++ b/api/core/tools/provider/builtin/audio/audio.yaml @@ -0,0 +1,11 @@ +identity: + author: hjlarry + name: audio + label: + en_US: Audio + description: + en_US: A tool for tts and asr. + zh_Hans: 一个用于文本转语音和语音转文本的工具。 + icon: icon.svg + tags: + - utilities diff --git a/api/core/tools/provider/builtin/audio/tools/asr.py b/api/core/tools/provider/builtin/audio/tools/asr.py new file mode 100644 index 0000000000000000000000000000000000000000..2aa409cd1ac082fc2638b50dc9064678a17fa7ae --- /dev/null +++ b/api/core/tools/provider/builtin/audio/tools/asr.py @@ -0,0 +1,69 @@ +import io +from typing import Any + +from core.file.enums import FileType +from core.file.file_manager import download +from core.model_manager import ModelManager +from core.model_runtime.entities.model_entities import ModelType +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter, ToolParameterOption +from core.tools.tool.builtin_tool import BuiltinTool +from services.model_provider_service import ModelProviderService + + +class ASRTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> list[ToolInvokeMessage]: + file = tool_parameters.get("audio_file") + if file.type != FileType.AUDIO: + return [self.create_text_message("not a valid audio file")] + audio_binary = io.BytesIO(download(file)) + audio_binary.name = "temp.mp3" + provider, model = tool_parameters.get("model").split("#") + model_manager = ModelManager() + model_instance = model_manager.get_model_instance( + tenant_id=self.runtime.tenant_id, + provider=provider, + model_type=ModelType.SPEECH2TEXT, + model=model, + ) + text = model_instance.invoke_speech2text( + file=audio_binary, + user=user_id, + ) + return [self.create_text_message(text)] + + def get_available_models(self) -> list[tuple[str, str]]: + model_provider_service = ModelProviderService() + models = model_provider_service.get_models_by_model_type( + tenant_id=self.runtime.tenant_id, model_type="speech2text" + ) + items = [] + for provider_model in models: + provider = provider_model.provider + for model in provider_model.models: + items.append((provider, model.model)) + return items + + def get_runtime_parameters(self) -> list[ToolParameter]: + parameters = [] + + options = [] + for provider, model in self.get_available_models(): + option = ToolParameterOption(value=f"{provider}#{model}", label=I18nObject(en_US=f"{model}({provider})")) + options.append(option) + + parameters.append( + ToolParameter( + name="model", + label=I18nObject(en_US="Model", zh_Hans="Model"), + human_description=I18nObject( + en_US="All available ASR models. You can config model in the Model Provider of Settings.", + zh_Hans="所有可用的 ASR 模型。你可以在设置中的模型供应商里配置。", + ), + type=ToolParameter.ToolParameterType.SELECT, + form=ToolParameter.ToolParameterForm.FORM, + required=True, + options=options, + ) + ) + return parameters diff --git a/api/core/tools/provider/builtin/audio/tools/asr.yaml b/api/core/tools/provider/builtin/audio/tools/asr.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b2c82f8086379dbe28d120734a1aab4facc5aa0c --- /dev/null +++ b/api/core/tools/provider/builtin/audio/tools/asr.yaml @@ -0,0 +1,22 @@ +identity: + name: asr + author: hjlarry + label: + en_US: Speech To Text +description: + human: + en_US: Convert audio file to text. + zh_Hans: 将音频文件转换为文本。 + llm: Convert audio file to text. +parameters: + - name: audio_file + type: file + required: true + label: + en_US: Audio File + zh_Hans: 音频文件 + human_description: + en_US: The audio file to be converted. + zh_Hans: 要转换的音频文件。 + llm_description: The audio file to be converted. + form: llm diff --git a/api/core/tools/provider/builtin/audio/tools/tts.py b/api/core/tools/provider/builtin/audio/tools/tts.py new file mode 100644 index 0000000000000000000000000000000000000000..8a33ac405bd4c3a88c6a1efca00154083778497c --- /dev/null +++ b/api/core/tools/provider/builtin/audio/tools/tts.py @@ -0,0 +1,97 @@ +import io +from typing import Any + +from core.model_manager import ModelManager +from core.model_runtime.entities.model_entities import ModelPropertyKey, ModelType +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter, ToolParameterOption +from core.tools.tool.builtin_tool import BuiltinTool +from services.model_provider_service import ModelProviderService + + +class TTSTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> list[ToolInvokeMessage]: + provider, model = tool_parameters.get("model", "").split("#") + voice = tool_parameters.get(f"voice#{provider}#{model}", "") + model_manager = ModelManager() + if not self.runtime: + raise ValueError("Runtime is required") + model_instance = model_manager.get_model_instance( + tenant_id=self.runtime.tenant_id or "", + provider=provider, + model_type=ModelType.TTS, + model=model, + ) + tts = model_instance.invoke_tts( + content_text=tool_parameters.get("text", ""), + user=user_id, + tenant_id=self.runtime.tenant_id or "", + voice=voice, + ) + buffer = io.BytesIO() + for chunk in tts: + buffer.write(chunk) + + wav_bytes = buffer.getvalue() + return [ + self.create_text_message("Audio generated successfully"), + self.create_blob_message( + blob=wav_bytes, + meta={"mime_type": "audio/x-wav"}, + save_as=self.VariableKey.AUDIO, + ), + ] + + def get_available_models(self) -> list[tuple[str, str, list[Any]]]: + if not self.runtime: + raise ValueError("Runtime is required") + model_provider_service = ModelProviderService() + tid: str = self.runtime.tenant_id or "" + models = model_provider_service.get_models_by_model_type(tenant_id=tid, model_type="tts") + items = [] + for provider_model in models: + provider = provider_model.provider + for model in provider_model.models: + voices = model.model_properties.get(ModelPropertyKey.VOICES, []) + items.append((provider, model.model, voices)) + return items + + def get_runtime_parameters(self) -> list[ToolParameter]: + parameters = [] + + options = [] + for provider, model, voices in self.get_available_models(): + option = ToolParameterOption(value=f"{provider}#{model}", label=I18nObject(en_US=f"{model}({provider})")) + options.append(option) + parameters.append( + ToolParameter( + name=f"voice#{provider}#{model}", + label=I18nObject(en_US=f"Voice of {model}({provider})"), + human_description=I18nObject(en_US=f"Select a voice for {model} model"), + placeholder=I18nObject(en_US="Select a voice"), + type=ToolParameter.ToolParameterType.SELECT, + form=ToolParameter.ToolParameterForm.FORM, + options=[ + ToolParameterOption(value=voice.get("mode"), label=I18nObject(en_US=voice.get("name"))) + for voice in voices + ], + ) + ) + + parameters.insert( + 0, + ToolParameter( + name="model", + label=I18nObject(en_US="Model", zh_Hans="Model"), + human_description=I18nObject( + en_US="All available TTS models. You can config model in the Model Provider of Settings.", + zh_Hans="所有可用的 TTS 模型。你可以在设置中的模型供应商里配置。", + ), + type=ToolParameter.ToolParameterType.SELECT, + form=ToolParameter.ToolParameterForm.FORM, + required=True, + placeholder=I18nObject(en_US="Select a model", zh_Hans="选择模型"), + options=options, + ), + ) + return parameters diff --git a/api/core/tools/provider/builtin/audio/tools/tts.yaml b/api/core/tools/provider/builtin/audio/tools/tts.yaml new file mode 100644 index 0000000000000000000000000000000000000000..36f42bd689fc7108c33acb5fd6bd33058455ded0 --- /dev/null +++ b/api/core/tools/provider/builtin/audio/tools/tts.yaml @@ -0,0 +1,22 @@ +identity: + name: tts + author: hjlarry + label: + en_US: Text To Speech +description: + human: + en_US: Convert text to audio file. + zh_Hans: 将文本转换为音频文件。 + llm: Convert text to audio file. +parameters: + - name: text + type: string + required: true + label: + en_US: Text + zh_Hans: 文本 + human_description: + en_US: The text to be converted. + zh_Hans: 要转换的文本。 + llm_description: The text to be converted. + form: llm diff --git a/api/core/tools/provider/builtin/aws/_assets/icon.svg b/api/core/tools/provider/builtin/aws/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..ecfcfc08d4eeff8fa10621d11791cadfd152edee --- /dev/null +++ b/api/core/tools/provider/builtin/aws/_assets/icon.svg @@ -0,0 +1,9 @@ + + + + + + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/aws/aws.py b/api/core/tools/provider/builtin/aws/aws.py new file mode 100644 index 0000000000000000000000000000000000000000..f81b5dbd27d17caba0ad40744d0995de1ac3b895 --- /dev/null +++ b/api/core/tools/provider/builtin/aws/aws.py @@ -0,0 +1,24 @@ +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.aws.tools.sagemaker_text_rerank import SageMakerReRankTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class SageMakerProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + try: + SageMakerReRankTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "sagemaker_endpoint": "", + "query": "misaka mikoto", + "candidate_texts": "hello$$$hello world", + "topk": 5, + "aws_region": "", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/aws/aws.yaml b/api/core/tools/provider/builtin/aws/aws.yaml new file mode 100644 index 0000000000000000000000000000000000000000..847c6824a53df65803d662a1632753c552494d34 --- /dev/null +++ b/api/core/tools/provider/builtin/aws/aws.yaml @@ -0,0 +1,15 @@ +identity: + author: AWS + name: aws + label: + en_US: AWS + zh_Hans: 亚马逊云科技 + pt_BR: AWS + description: + en_US: Services on AWS. + zh_Hans: 亚马逊云科技的各类服务 + pt_BR: Services on AWS. + icon: icon.svg + tags: + - search +credentials_for_provider: diff --git a/api/core/tools/provider/builtin/aws/tools/apply_guardrail.py b/api/core/tools/provider/builtin/aws/tools/apply_guardrail.py new file mode 100644 index 0000000000000000000000000000000000000000..b224ff5258c8791670113bdc64a2b036fab65fba --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/apply_guardrail.py @@ -0,0 +1,90 @@ +import json +import logging +from typing import Any, Union + +import boto3 # type: ignore +from botocore.exceptions import BotoCoreError # type: ignore +from pydantic import BaseModel, Field + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + + +class GuardrailParameters(BaseModel): + guardrail_id: str = Field(..., description="The identifier of the guardrail") + guardrail_version: str = Field(..., description="The version of the guardrail") + source: str = Field(..., description="The source of the content") + text: str = Field(..., description="The text to apply the guardrail to") + aws_region: str = Field(..., description="AWS region for the Bedrock client") + + +class ApplyGuardrailTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Invoke the ApplyGuardrail tool + """ + try: + # Validate and parse input parameters + params = GuardrailParameters(**tool_parameters) + + # Initialize AWS client + bedrock_client = boto3.client("bedrock-runtime", region_name=params.aws_region) + + # Apply guardrail + response = bedrock_client.apply_guardrail( + guardrailIdentifier=params.guardrail_id, + guardrailVersion=params.guardrail_version, + source=params.source, + content=[{"text": {"text": params.text}}], + ) + + logger.info(f"Raw response from AWS: {json.dumps(response, indent=2)}") + + # Check for empty response + if not response: + return self.create_text_message(text="Received empty response from AWS Bedrock.") + + # Process the result + action = response.get("action", "No action specified") + outputs = response.get("outputs", []) + output = outputs[0].get("text", "No output received") if outputs else "No output received" + assessments = response.get("assessments", []) + + # Format assessments + formatted_assessments = [] + for assessment in assessments: + for policy_type, policy_data in assessment.items(): + if isinstance(policy_data, dict) and "topics" in policy_data: + for topic in policy_data["topics"]: + formatted_assessments.append( + f"Policy: {policy_type}, Topic: {topic['name']}, Type: {topic['type']}," + f" Action: {topic['action']}" + ) + else: + formatted_assessments.append(f"Policy: {policy_type}, Data: {policy_data}") + + result = f"Action: {action}\n " + result += f"Output: {output}\n " + if formatted_assessments: + result += "Assessments:\n " + "\n ".join(formatted_assessments) + "\n " + # result += f"Full response: {json.dumps(response, indent=2, ensure_ascii=False)}" + + return self.create_text_message(text=result) + + except BotoCoreError as e: + error_message = f"AWS service error: {str(e)}" + logger.error(error_message, exc_info=True) + return self.create_text_message(text=error_message) + except json.JSONDecodeError as e: + error_message = f"JSON parsing error: {str(e)}" + logger.error(error_message, exc_info=True) + return self.create_text_message(text=error_message) + except Exception as e: + error_message = f"An unexpected error occurred: {str(e)}" + logger.error(error_message, exc_info=True) + return self.create_text_message(text=error_message) diff --git a/api/core/tools/provider/builtin/aws/tools/apply_guardrail.yaml b/api/core/tools/provider/builtin/aws/tools/apply_guardrail.yaml new file mode 100644 index 0000000000000000000000000000000000000000..66044e4ea84fe15c481e5040d4cbfce85fe9e187 --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/apply_guardrail.yaml @@ -0,0 +1,67 @@ +identity: + name: apply_guardrail + author: AWS + label: + en_US: Content Moderation Guardrails + zh_Hans: 内容审查护栏 +description: + human: + en_US: Content Moderation Guardrails utilizes the ApplyGuardrail API, a feature of Guardrails for Amazon Bedrock. This API is capable of evaluating input prompts and model responses for all Foundation Models (FMs), including those on Amazon Bedrock, custom FMs, and third-party FMs. By implementing this functionality, organizations can achieve centralized governance across all their generative AI applications, thereby enhancing control and consistency in content moderation. + zh_Hans: 内容审查护栏采用 Guardrails for Amazon Bedrock 功能中的 ApplyGuardrail API 。ApplyGuardrail 可以评估所有基础模型(FMs)的输入提示和模型响应,包括 Amazon Bedrock 上的 FMs、自定义 FMs 和第三方 FMs。通过实施这一功能, 组织可以在所有生成式 AI 应用程序中实现集中化的治理,从而增强内容审核的控制力和一致性。 + llm: Content Moderation Guardrails utilizes the ApplyGuardrail API, a feature of Guardrails for Amazon Bedrock. This API is capable of evaluating input prompts and model responses for all Foundation Models (FMs), including those on Amazon Bedrock, custom FMs, and third-party FMs. By implementing this functionality, organizations can achieve centralized governance across all their generative AI applications, thereby enhancing control and consistency in content moderation. +parameters: + - name: guardrail_id + type: string + required: true + label: + en_US: Guardrail ID + zh_Hans: Guardrail ID + human_description: + en_US: Please enter the ID of the Guardrail that has already been created on Amazon Bedrock, for example 'qk5nk0e4b77b'. + zh_Hans: 请输入已经在 Amazon Bedrock 上创建好的 Guardrail ID, 例如 'qk5nk0e4b77b'. + llm_description: Please enter the ID of the Guardrail that has already been created on Amazon Bedrock, for example 'qk5nk0e4b77b'. + form: form + - name: guardrail_version + type: string + required: true + label: + en_US: Guardrail Version Number + zh_Hans: Guardrail 版本号码 + human_description: + en_US: Please enter the published version of the Guardrail ID that has already been created on Amazon Bedrock. This is typically a version number, such as 2. + zh_Hans: 请输入已经在Amazon Bedrock 上创建好的Guardrail ID发布的版本, 通常使用版本号, 例如2. + llm_description: Please enter the published version of the Guardrail ID that has already been created on Amazon Bedrock. This is typically a version number, such as 2. + form: form + - name: source + type: string + required: true + label: + en_US: Content Source (INPUT or OUTPUT) + zh_Hans: 内容来源 (INPUT or OUTPUT) + human_description: + en_US: The source of data used in the request to apply the guardrail. Valid Values "INPUT | OUTPUT" + zh_Hans: 用于应用护栏的请求中所使用的数据来源。有效值为 "INPUT | OUTPUT" + llm_description: The source of data used in the request to apply the guardrail. Valid Values "INPUT | OUTPUT" + form: form + - name: text + type: string + required: true + label: + en_US: Content to be reviewed + zh_Hans: 待审查内容 + human_description: + en_US: The content used for requesting guardrail review, which can be either user input or LLM output. + zh_Hans: 用于请求护栏审查的内容,可以是用户输入或 LLM 输出。 + llm_description: The content used for requesting guardrail review, which can be either user input or LLM output. + form: llm + - name: aws_region + type: string + required: true + label: + en_US: AWS Region + zh_Hans: AWS 区域 + human_description: + en_US: Please enter the AWS region for the Bedrock client, for example 'us-east-1'. + zh_Hans: 请输入 Bedrock 客户端的 AWS 区域,例如 'us-east-1'。 + llm_description: Please enter the AWS region for the Bedrock client, for example 'us-east-1'. + form: form diff --git a/api/core/tools/provider/builtin/aws/tools/bedrock_retrieve.py b/api/core/tools/provider/builtin/aws/tools/bedrock_retrieve.py new file mode 100644 index 0000000000000000000000000000000000000000..19e7bfa76eb844350d1511229fa2258ca95be6c9 --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/bedrock_retrieve.py @@ -0,0 +1,162 @@ +import json +import operator +from typing import Any, Optional, Union + +import boto3 + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class BedrockRetrieveTool(BuiltinTool): + bedrock_client: Any = None + knowledge_base_id: str = None + topk: int = None + + def _bedrock_retrieve( + self, + query_input: str, + knowledge_base_id: str, + num_results: int, + search_type: str, + rerank_model_id: str, + metadata_filter: Optional[dict] = None, + ): + try: + retrieval_query = {"text": query_input} + + if search_type not in ["HYBRID", "SEMANTIC"]: + raise RuntimeException("search_type should be HYBRID or SEMANTIC") + + retrieval_configuration = { + "vectorSearchConfiguration": {"numberOfResults": num_results, "overrideSearchType": search_type} + } + + if rerank_model_id != "default": + model_for_rerank_arn = f"arn:aws:bedrock:us-west-2::foundation-model/{rerank_model_id}" + rerankingConfiguration = { + "bedrockRerankingConfiguration": { + "numberOfRerankedResults": num_results, + "modelConfiguration": {"modelArn": model_for_rerank_arn}, + }, + "type": "BEDROCK_RERANKING_MODEL", + } + + retrieval_configuration["vectorSearchConfiguration"]["rerankingConfiguration"] = rerankingConfiguration + retrieval_configuration["vectorSearchConfiguration"]["numberOfResults"] = num_results * 5 + + # 如果有元数据过滤条件,则添加到检索配置中 + if metadata_filter: + retrieval_configuration["vectorSearchConfiguration"]["filter"] = metadata_filter + + response = self.bedrock_client.retrieve( + knowledgeBaseId=knowledge_base_id, + retrievalQuery=retrieval_query, + retrievalConfiguration=retrieval_configuration, + ) + + results = [] + for result in response.get("retrievalResults", []): + results.append( + { + "content": result.get("content", {}).get("text", ""), + "score": result.get("score", 0.0), + "metadata": result.get("metadata", {}), + } + ) + + return results + except Exception as e: + raise Exception(f"Error retrieving from knowledge base: {str(e)}") + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + try: + line = 0 + # Initialize Bedrock client if not already initialized + if not self.bedrock_client: + aws_region = tool_parameters.get("aws_region") + aws_access_key_id = tool_parameters.get("aws_access_key_id") + aws_secret_access_key = tool_parameters.get("aws_secret_access_key") + + client_kwargs = {"service_name": "bedrock-agent-runtime", "region_name": aws_region or None} + + # Only add credentials if both access key and secret key are provided + if aws_access_key_id and aws_secret_access_key: + client_kwargs.update( + {"aws_access_key_id": aws_access_key_id, "aws_secret_access_key": aws_secret_access_key} + ) + + self.bedrock_client = boto3.client(**client_kwargs) + except Exception as e: + return self.create_text_message(f"Failed to initialize Bedrock client: {str(e)}") + + try: + line = 1 + if not self.knowledge_base_id: + self.knowledge_base_id = tool_parameters.get("knowledge_base_id") + if not self.knowledge_base_id: + return self.create_text_message("Please provide knowledge_base_id") + + line = 2 + if not self.topk: + self.topk = tool_parameters.get("topk", 5) + + line = 3 + query = tool_parameters.get("query", "") + if not query: + return self.create_text_message("Please input query") + + # 获取元数据过滤条件(如果存在) + metadata_filter_str = tool_parameters.get("metadata_filter") + metadata_filter = json.loads(metadata_filter_str) if metadata_filter_str else None + + search_type = tool_parameters.get("search_type") + rerank_model_id = tool_parameters.get("rerank_model_id") + + line = 4 + retrieved_docs = self._bedrock_retrieve( + query_input=query, + knowledge_base_id=self.knowledge_base_id, + num_results=self.topk, + search_type=search_type, + rerank_model_id=rerank_model_id, + metadata_filter=metadata_filter, + ) + + line = 5 + # Sort results by score in descending order + sorted_docs = sorted(retrieved_docs, key=operator.itemgetter("score"), reverse=True) + + line = 6 + result_type = tool_parameters.get("result_type") + if result_type == "json": + return [self.create_json_message(res) for res in sorted_docs] + else: + text = "" + for i, res in enumerate(sorted_docs): + text += f"{i + 1}: {res['content']}\n" + return self.create_text_message(text) + + except Exception as e: + return self.create_text_message(f"Exception {str(e)}, line : {line}") + + def validate_parameters(self, parameters: dict[str, Any]) -> None: + """ + Validate the parameters + """ + if not parameters.get("knowledge_base_id"): + raise ValueError("knowledge_base_id is required") + + if not parameters.get("query"): + raise ValueError("query is required") + + metadata_filter_str = parameters.get("metadata_filter") + if metadata_filter_str and not isinstance(json.loads(metadata_filter_str), dict): + raise ValueError("metadata_filter must be a valid JSON object") diff --git a/api/core/tools/provider/builtin/aws/tools/bedrock_retrieve.yaml b/api/core/tools/provider/builtin/aws/tools/bedrock_retrieve.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d0c520b39298d15eaf59dcce8ae68cc8611f8890 --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/bedrock_retrieve.yaml @@ -0,0 +1,179 @@ +identity: + name: bedrock_retrieve + author: AWS + label: + en_US: Bedrock Retrieve + zh_Hans: Bedrock检索 + pt_BR: Bedrock Retrieve + icon: icon.svg + +description: + human: + en_US: A tool for retrieving relevant information from Amazon Bedrock Knowledge Base. You can find deploy instructions on Github Repo - https://github.com/aws-samples/dify-aws-tool + zh_Hans: Amazon Bedrock知识库检索工具, 请参考 Github Repo - https://github.com/aws-samples/dify-aws-tool上的部署说明 + pt_BR: A tool for retrieving relevant information from Amazon Bedrock Knowledge Base. + llm: A tool for retrieving relevant information from Amazon Bedrock Knowledge Base. You can find deploy instructions on Github Repo - https://github.com/aws-samples/dify-aws-tool + +parameters: + - name: aws_region + type: string + required: false + label: + en_US: AWS Region + zh_Hans: AWS区域 + human_description: + en_US: AWS region for the Bedrock service + zh_Hans: Bedrock服务的AWS区域 + form: form + + - name: aws_access_key_id + type: string + required: false + label: + en_US: AWS Access Key ID + zh_Hans: AWS访问密钥ID + human_description: + en_US: AWS access key ID for authentication (optional) + zh_Hans: 用于身份验证的AWS访问密钥ID(可选) + form: form + + - name: aws_secret_access_key + type: string + required: false + label: + en_US: AWS Secret Access Key + zh_Hans: AWS秘密访问密钥 + human_description: + en_US: AWS secret access key for authentication (optional) + zh_Hans: 用于身份验证的AWS秘密访问密钥(可选) + form: form + + - name: result_type + type: select + required: true + label: + en_US: result type + zh_Hans: 结果类型 + human_description: + en_US: return a list of json or texts + zh_Hans: 返回一个列表,内容是json还是纯文本 + default: text + options: + - value: json + label: + en_US: JSON + zh_Hans: JSON + - value: text + label: + en_US: Text + zh_Hans: 文本 + form: form + + - name: knowledge_base_id + type: string + required: true + label: + en_US: Bedrock Knowledge Base ID + zh_Hans: Bedrock知识库ID + pt_BR: Bedrock Knowledge Base ID + human_description: + en_US: ID of the Bedrock Knowledge Base to retrieve from + zh_Hans: 用于检索的Bedrock知识库ID + pt_BR: ID of the Bedrock Knowledge Base to retrieve from + llm_description: ID of the Bedrock Knowledge Base to retrieve from + form: form + + - name: query + type: string + required: true + label: + en_US: Query string + zh_Hans: 查询语句 + pt_BR: Query string + human_description: + en_US: The search query to retrieve relevant information + zh_Hans: 用于检索相关信息的查询语句 + pt_BR: The search query to retrieve relevant information + llm_description: The search query to retrieve relevant information + form: llm + + - name: topk + type: number + required: false + form: form + label: + en_US: Limit for results count + zh_Hans: 返回结果数量限制 + pt_BR: Limit for results count + human_description: + en_US: Maximum number of results to return + zh_Hans: 最大返回结果数量 + pt_BR: Maximum number of results to return + min: 1 + max: 10 + default: 5 + + - name: search_type + type: select + required: false + label: + en_US: search type + zh_Hans: 搜索类型 + pt_BR: search type + human_description: + en_US: search type + zh_Hans: 搜索类型 + pt_BR: search type + llm_description: search type + default: SEMANTIC + options: + - value: SEMANTIC + label: + en_US: SEMANTIC + zh_Hans: 语义搜索 + - value: HYBRID + label: + en_US: HYBRID + zh_Hans: 混合搜索 + form: form + + - name: rerank_model_id + type: select + required: false + label: + en_US: rerank model id + zh_Hans: 重拍模型ID + pt_BR: rerank model id + human_description: + en_US: rerank model id + zh_Hans: 重拍模型ID + pt_BR: rerank model id + llm_description: rerank model id + default: default + options: + - value: default + label: + en_US: default + zh_Hans: 默认 + - value: cohere.rerank-v3-5:0 + label: + en_US: cohere.rerank-v3-5:0 + zh_Hans: cohere.rerank-v3-5:0 + - value: amazon.rerank-v1:0 + label: + en_US: amazon.rerank-v1:0 + zh_Hans: amazon.rerank-v1:0 + form: form + + - name: metadata_filter # Additional parameter for metadata filtering + type: string # String type, expects JSON-formatted filter conditions + required: false # Optional field - can be omitted + label: + en_US: Metadata Filter + zh_Hans: 元数据过滤器 + pt_BR: Metadata Filter + human_description: + en_US: 'JSON formatted filter conditions for metadata (e.g., {"greaterThan": {"key: "aaa", "value": 10}})' + zh_Hans: '元数据的JSON格式过滤条件(例如,{{"greaterThan": {"key: "aaa", "value": 10}})' + pt_BR: 'JSON formatted filter conditions for metadata (e.g., {"greaterThan": {"key: "aaa", "value": 10}})' + form: form diff --git a/api/core/tools/provider/builtin/aws/tools/bedrock_retrieve_and_generate.py b/api/core/tools/provider/builtin/aws/tools/bedrock_retrieve_and_generate.py new file mode 100644 index 0000000000000000000000000000000000000000..3717aac344fe5d7f7afabf4635f2d65801493169 --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/bedrock_retrieve_and_generate.py @@ -0,0 +1,137 @@ +import json +from typing import Any + +import boto3 + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class BedrockRetrieveAndGenerateTool(BuiltinTool): + bedrock_client: Any = None + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> ToolInvokeMessage: + try: + # Initialize Bedrock client if not already initialized + if not self.bedrock_client: + aws_region = tool_parameters.get("aws_region") + aws_access_key_id = tool_parameters.get("aws_access_key_id") + aws_secret_access_key = tool_parameters.get("aws_secret_access_key") + + client_kwargs = {"service_name": "bedrock-agent-runtime", "region_name": aws_region or None} + + # Only add credentials if both access key and secret key are provided + if aws_access_key_id and aws_secret_access_key: + client_kwargs.update( + {"aws_access_key_id": aws_access_key_id, "aws_secret_access_key": aws_secret_access_key} + ) + + self.bedrock_client = boto3.client(**client_kwargs) + except Exception as e: + return self.create_text_message(f"Failed to initialize Bedrock client: {str(e)}") + + try: + request_config = {} + + # Set input configuration + input_text = tool_parameters.get("input") + if input_text: + request_config["input"] = {"text": input_text} + + # Build retrieve and generate configuration + config_type = tool_parameters.get("type") + retrieve_generate_config = {"type": config_type} + + # Add configuration based on type + if config_type == "KNOWLEDGE_BASE": + kb_config_str = tool_parameters.get("knowledge_base_configuration") + kb_config = json.loads(kb_config_str) if kb_config_str else None + retrieve_generate_config["knowledgeBaseConfiguration"] = kb_config + else: # EXTERNAL_SOURCES + es_config_str = tool_parameters.get("external_sources_configuration") + es_config = json.loads(kb_config_str) if es_config_str else None + retrieve_generate_config["externalSourcesConfiguration"] = es_config + + request_config["retrieveAndGenerateConfiguration"] = retrieve_generate_config + + # Parse session configuration + session_config_str = tool_parameters.get("session_configuration") + session_config = json.loads(session_config_str) if session_config_str else None + if session_config: + request_config["sessionConfiguration"] = session_config + + # Add session ID if provided + session_id = tool_parameters.get("session_id") + if session_id: + request_config["sessionId"] = session_id + + # Send request + response = self.bedrock_client.retrieve_and_generate(**request_config) + + # Process response + result = {"output": response.get("output", {}).get("text", ""), "citations": []} + + # Process citations + for citation in response.get("citations", []): + citation_info = { + "text": citation.get("generatedResponsePart", {}).get("textResponsePart", {}).get("text", ""), + "references": [], + } + + for ref in citation.get("retrievedReferences", []): + reference = { + "content": ref.get("content", {}).get("text", ""), + "metadata": ref.get("metadata", {}), + "location": None, + } + + location = ref.get("location", {}) + if location.get("type") == "S3": + reference["location"] = location.get("s3Location", {}).get("uri") + + citation_info["references"].append(reference) + + result["citations"].append(citation_info) + result_type = tool_parameters.get("result_type") + if result_type == "json": + return self.create_json_message(result) + elif result_type == "text-with-citations": + return self.create_text_message(result) + else: + return self.create_text_message(result.get("output")) + except json.JSONDecodeError as e: + return self.create_text_message(f"Invalid JSON format: {str(e)}") + except Exception as e: + return self.create_text_message(f"Tool invocation error: {str(e)}") + + def validate_parameters(self, parameters: dict[str, Any]) -> None: + """Validate the parameters""" + # Validate required parameters + if not parameters.get("input"): + raise ValueError("input is required") + if not parameters.get("type"): + raise ValueError("type is required") + + # Validate JSON configurations + json_configs = ["knowledge_base_configuration", "external_sources_configuration", "session_configuration"] + for config in json_configs: + if config_value := parameters.get(config): + try: + json.loads(config_value) + except json.JSONDecodeError: + raise ValueError(f"{config} must be a valid JSON string") + + # Validate configuration type + config_type = parameters.get("type") + if config_type not in ["KNOWLEDGE_BASE", "EXTERNAL_SOURCES"]: + raise ValueError("type must be either KNOWLEDGE_BASE or EXTERNAL_SOURCES") + + # Validate type-specific configuration + if config_type == "KNOWLEDGE_BASE" and not parameters.get("knowledge_base_configuration"): + raise ValueError("knowledge_base_configuration is required when type is KNOWLEDGE_BASE") + elif config_type == "EXTERNAL_SOURCES" and not parameters.get("external_sources_configuration"): + raise ValueError("external_sources_configuration is required when type is EXTERNAL_SOURCES") diff --git a/api/core/tools/provider/builtin/aws/tools/bedrock_retrieve_and_generate.yaml b/api/core/tools/provider/builtin/aws/tools/bedrock_retrieve_and_generate.yaml new file mode 100644 index 0000000000000000000000000000000000000000..68f418fc5caa7678262e8003dcbfd2e3b120e629 --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/bedrock_retrieve_and_generate.yaml @@ -0,0 +1,148 @@ +identity: + name: bedrock_retrieve_and_generate + author: AWS + label: + en_US: Bedrock Retrieve and Generate + zh_Hans: Bedrock检索和生成 + icon: icon.svg + +description: + human: + en_US: "This is an advanced usage of Bedrock Retrieve. Please refer to the API documentation for detailed parameters and paste them into the corresponding Knowledge Base Configuration or External Sources Configuration" + zh_Hans: "这个工具为Bedrock Retrieve的高级用法,请参考API设置详细的参数,并粘贴到对应的知识库配置或者外部源配置" + llm: A tool for retrieving and generating information using Amazon Bedrock Knowledge Base + +parameters: + - name: aws_region + type: string + required: false + label: + en_US: AWS Region + zh_Hans: AWS区域 + human_description: + en_US: AWS region for the Bedrock service + zh_Hans: Bedrock服务的AWS区域 + form: form + + - name: aws_access_key_id + type: string + required: false + label: + en_US: AWS Access Key ID + zh_Hans: AWS访问密钥ID + human_description: + en_US: AWS access key ID for authentication (optional) + zh_Hans: 用于身份验证的AWS访问密钥ID(可选) + form: form + + - name: aws_secret_access_key + type: string + required: false + label: + en_US: AWS Secret Access Key + zh_Hans: AWS秘密访问密钥 + human_description: + en_US: AWS secret access key for authentication (optional) + zh_Hans: 用于身份验证的AWS秘密访问密钥(可选) + form: form + + - name: result_type + type: select + required: true + label: + en_US: result type + zh_Hans: 结果类型 + human_description: + en_US: return a list of json or texts + zh_Hans: 返回一个列表,内容是json还是纯文本 + default: text + options: + - value: json + label: + en_US: JSON + zh_Hans: JSON + - value: text + label: + en_US: Text + zh_Hans: 文本 + - value: text-with-citations + label: + en_US: Text With Citations + zh_Hans: 文本(包含引用) + form: form + + - name: input + type: string + required: true + label: + en_US: Input Text + zh_Hans: 输入文本 + human_description: + en_US: The text query to retrieve information + zh_Hans: 用于检索信息的文本查询 + form: llm + + - name: type + type: select + required: true + label: + en_US: Configuration Type + zh_Hans: 配置类型 + human_description: + en_US: Type of retrieve and generate configuration + zh_Hans: 检索和生成配置的类型 + options: + - value: KNOWLEDGE_BASE + label: + en_US: Knowledge Base + zh_Hans: 知识库 + - value: EXTERNAL_SOURCES + label: + en_US: External Sources + zh_Hans: 外部源 + form: form + + - name: knowledge_base_configuration + type: string + required: false + label: + en_US: Knowledge Base Configuration + zh_Hans: 知识库配置 + human_description: + en_US: Please refer to @https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent-runtime/client/retrieve_and_generate.html#retrieve-and-generate for complete parameters and paste them here + zh_Hans: 请参考 https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent-runtime/client/retrieve_and_generate.html#retrieve-and-generate 配置完整的参数并粘贴到这里 + form: form + + - name: external_sources_configuration + type: string + required: false + label: + en_US: External Sources Configuration + zh_Hans: 外部源配置 + human_description: + en_US: Please refer to https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent-runtime/client/retrieve_and_generate.html#retrieve-and-generate for complete parameters and paste them here + zh_Hans: 请参考 https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent-runtime/client/retrieve_and_generate.html#retrieve-and-generate 配置完整的参数并粘贴到这里 + form: form + + - name: session_configuration + type: string + required: false + label: + en_US: Session Configuration + zh_Hans: 会话配置 + human_description: + en_US: JSON formatted session configuration + zh_Hans: JSON格式的会话配置 + default: "" + form: form + + - name: session_id + type: string + required: false + label: + en_US: Session ID + zh_Hans: 会话ID + human_description: + en_US: Session ID for continuous conversations + zh_Hans: 用于连续对话的会话ID + form: form diff --git a/api/core/tools/provider/builtin/aws/tools/lambda_translate_utils.py b/api/core/tools/provider/builtin/aws/tools/lambda_translate_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b6d16d2759c30e8140fadfe939800fa76b2237dc --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/lambda_translate_utils.py @@ -0,0 +1,91 @@ +import json +from typing import Any, Union + +import boto3 # type: ignore + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class LambdaTranslateUtilsTool(BuiltinTool): + lambda_client: Any = None + + def _invoke_lambda(self, text_content, src_lang, dest_lang, model_id, dictionary_name, request_type, lambda_name): + msg = { + "src_contents": [text_content], + "src_lang": src_lang, + "dest_lang": dest_lang, + "dictionary_id": dictionary_name, + "request_type": request_type, + "model_id": model_id, + } + + invoke_response = self.lambda_client.invoke( + FunctionName=lambda_name, InvocationType="RequestResponse", Payload=json.dumps(msg) + ) + response_body = invoke_response["Payload"] + + response_str = response_body.read().decode("unicode_escape") + + return response_str + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + line = 0 + try: + if not self.lambda_client: + aws_region = tool_parameters.get("aws_region") + if aws_region: + self.lambda_client = boto3.client("lambda", region_name=aws_region) + else: + self.lambda_client = boto3.client("lambda") + + line = 1 + text_content = tool_parameters.get("text_content", "") + if not text_content: + return self.create_text_message("Please input text_content") + + line = 2 + src_lang = tool_parameters.get("src_lang", "") + if not src_lang: + return self.create_text_message("Please input src_lang") + + line = 3 + dest_lang = tool_parameters.get("dest_lang", "") + if not dest_lang: + return self.create_text_message("Please input dest_lang") + + line = 4 + lambda_name = tool_parameters.get("lambda_name", "") + if not lambda_name: + return self.create_text_message("Please input lambda_name") + + line = 5 + request_type = tool_parameters.get("request_type", "") + if not request_type: + return self.create_text_message("Please input request_type") + + line = 6 + model_id = tool_parameters.get("model_id", "") + if not model_id: + return self.create_text_message("Please input model_id") + + line = 7 + dictionary_name = tool_parameters.get("dictionary_name", "") + if not dictionary_name: + return self.create_text_message("Please input dictionary_name") + + result = self._invoke_lambda( + text_content, src_lang, dest_lang, model_id, dictionary_name, request_type, lambda_name + ) + + return self.create_text_message(text=result) + + except Exception as e: + return self.create_text_message(f"Exception {str(e)}, line : {line}") diff --git a/api/core/tools/provider/builtin/aws/tools/lambda_translate_utils.yaml b/api/core/tools/provider/builtin/aws/tools/lambda_translate_utils.yaml new file mode 100644 index 0000000000000000000000000000000000000000..646602fcd6c245c29eb5c408c39426e178ef5b5d --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/lambda_translate_utils.yaml @@ -0,0 +1,134 @@ +identity: + name: lambda_translate_utils + author: AWS + label: + en_US: TranslateTool + zh_Hans: 翻译工具 + pt_BR: TranslateTool + icon: icon.svg +description: + human: + en_US: A util tools for LLM translation, extra deployment is needed on AWS. Please refer Github Repo - https://github.com/aws-samples/rag-based-translation-with-dynamodb-and-bedrock + zh_Hans: 大语言模型翻译工具(专词映射获取),需要在AWS上进行额外部署,可参考Github Repo - https://github.com/aws-samples/rag-based-translation-with-dynamodb-and-bedrock + pt_BR: A util tools for LLM translation, specific Lambda Function deployment is needed on AWS. Please refer Github Repo - https://github.com/aws-samples/rag-based-translation-with-dynamodb-and-bedrock + llm: A util tools for translation. +parameters: + - name: text_content + type: string + required: true + label: + en_US: source content for translation + zh_Hans: 待翻译原文 + pt_BR: source content for translation + human_description: + en_US: source content for translation + zh_Hans: 待翻译原文 + pt_BR: source content for translation + llm_description: source content for translation + form: llm + - name: src_lang + type: string + required: true + label: + en_US: source language code + zh_Hans: 原文语言代号 + pt_BR: source language code + human_description: + en_US: source language code + zh_Hans: 原文语言代号 + pt_BR: source language code + llm_description: source language code + form: llm + - name: dest_lang + type: string + required: true + label: + en_US: target language code + zh_Hans: 目标语言代号 + pt_BR: target language code + human_description: + en_US: target language code + zh_Hans: 目标语言代号 + pt_BR: target language code + llm_description: target language code + form: llm + - name: aws_region + type: string + required: false + label: + en_US: region of Lambda + zh_Hans: Lambda 所在的region + pt_BR: region of Lambda + human_description: + en_US: region of Lambda + zh_Hans: Lambda 所在的region + pt_BR: region of Lambda + llm_description: region of Lambda + form: form + - name: model_id + type: string + required: false + default: anthropic.claude-3-sonnet-20240229-v1:0 + label: + en_US: LLM model_id in bedrock + zh_Hans: bedrock上的大语言模型model_id + pt_BR: LLM model_id in bedrock + human_description: + en_US: LLM model_id in bedrock + zh_Hans: bedrock上的大语言模型model_id + pt_BR: LLM model_id in bedrock + llm_description: LLM model_id in bedrock + form: form + - name: dictionary_name + type: string + required: false + label: + en_US: dictionary name for term mapping + zh_Hans: 专词映射表名称 + pt_BR: dictionary name for term mapping + human_description: + en_US: dictionary name for term mapping + zh_Hans: 专词映射表名称 + pt_BR: dictionary name for term mapping + llm_description: dictionary name for term mapping + form: form + - name: request_type + type: select + required: false + label: + en_US: request type + zh_Hans: 请求类型 + pt_BR: request type + human_description: + en_US: request type + zh_Hans: 请求类型 + pt_BR: request type + default: term_mapping + options: + - value: term_mapping + label: + en_US: term_mapping + zh_Hans: 专词映射 + - value: segment_only + label: + en_US: segment_only + zh_Hans: 仅切词 + - value: translate + label: + en_US: translate + zh_Hans: 翻译内容 + form: form + - name: lambda_name + type: string + default: "translate_tool" + required: true + label: + en_US: AWS Lambda for term mapping retrieval + zh_Hans: 专词召回映射 - AWS Lambda + pt_BR: lambda name for term mapping retrieval + human_description: + en_US: AWS Lambda for term mapping retrieval + zh_Hans: 专词召回映射 - AWS Lambda + pt_BR: AWS Lambda for term mapping retrieval + llm_description: AWS Lambda for term mapping retrieval + form: form diff --git a/api/core/tools/provider/builtin/aws/tools/lambda_yaml_to_json.py b/api/core/tools/provider/builtin/aws/tools/lambda_yaml_to_json.py new file mode 100644 index 0000000000000000000000000000000000000000..01bc596346c231efd85b3700c3b7bfd2faceaea6 --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/lambda_yaml_to_json.py @@ -0,0 +1,70 @@ +import json +import logging +from typing import Any, Union + +import boto3 # type: ignore + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +console_handler = logging.StreamHandler() +logger.addHandler(console_handler) + + +class LambdaYamlToJsonTool(BuiltinTool): + lambda_client: Any = None + + def _invoke_lambda(self, lambda_name: str, yaml_content: str) -> str: + msg = {"body": yaml_content} + logger.info(json.dumps(msg)) + + invoke_response = self.lambda_client.invoke( + FunctionName=lambda_name, InvocationType="RequestResponse", Payload=json.dumps(msg) + ) + response_body = invoke_response["Payload"] + + response_str = response_body.read().decode("utf-8") + resp_json = json.loads(response_str) + + logger.info(resp_json) + if resp_json["statusCode"] != 200: + raise Exception(f"Invalid status code: {response_str}") + + return resp_json["body"] + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + try: + if not self.lambda_client: + aws_region = tool_parameters.get("aws_region") # todo: move aws_region out, and update client region + if aws_region: + self.lambda_client = boto3.client("lambda", region_name=aws_region) + else: + self.lambda_client = boto3.client("lambda") + + yaml_content = tool_parameters.get("yaml_content", "") + if not yaml_content: + return self.create_text_message("Please input yaml_content") + + lambda_name = tool_parameters.get("lambda_name", "") + if not lambda_name: + return self.create_text_message("Please input lambda_name") + logger.debug(f"{json.dumps(tool_parameters, indent=2, ensure_ascii=False)}") + + result = self._invoke_lambda(lambda_name, yaml_content) + logger.debug(result) + + return self.create_text_message(result) + except Exception as e: + return self.create_text_message(f"Exception: {str(e)}") + + console_handler.flush() diff --git a/api/core/tools/provider/builtin/aws/tools/lambda_yaml_to_json.yaml b/api/core/tools/provider/builtin/aws/tools/lambda_yaml_to_json.yaml new file mode 100644 index 0000000000000000000000000000000000000000..919c285348df83710159853a0d5dfae7036a2dec --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/lambda_yaml_to_json.yaml @@ -0,0 +1,53 @@ +identity: + name: lambda_yaml_to_json + author: AWS + label: + en_US: LambdaYamlToJson + zh_Hans: LambdaYamlToJson + pt_BR: LambdaYamlToJson + icon: icon.svg +description: + human: + en_US: A tool to convert yaml to json using AWS Lambda. + zh_Hans: 将 YAML 转为 JSON 的工具(通过AWS Lambda)。 + pt_BR: A tool to convert yaml to json using AWS Lambda. + llm: A tool to convert yaml to json. +parameters: + - name: yaml_content + type: string + required: true + label: + en_US: YAML content to convert for + zh_Hans: YAML 内容 + pt_BR: YAML content to convert for + human_description: + en_US: YAML content to convert for + zh_Hans: YAML 内容 + pt_BR: YAML content to convert for + llm_description: YAML content to convert for + form: llm + - name: aws_region + type: string + required: false + label: + en_US: region of lambda + zh_Hans: Lambda 所在的region + pt_BR: region of lambda + human_description: + en_US: region of lambda + zh_Hans: Lambda 所在的region + pt_BR: region of lambda + llm_description: region of lambda + form: form + - name: lambda_name + type: string + required: false + label: + en_US: name of lambda + zh_Hans: Lambda 名称 + pt_BR: name of lambda + human_description: + en_US: name of lambda + zh_Hans: Lambda 名称 + pt_BR: name of lambda + form: form diff --git a/api/core/tools/provider/builtin/aws/tools/nova_canvas.py b/api/core/tools/provider/builtin/aws/tools/nova_canvas.py new file mode 100644 index 0000000000000000000000000000000000000000..954dbe35a4a784fea53d5ba021f69168b2fb3095 --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/nova_canvas.py @@ -0,0 +1,357 @@ +import base64 +import json +import logging +import re +from datetime import datetime +from typing import Any, Union +from urllib.parse import urlparse + +import boto3 + +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter +from core.tools.tool.builtin_tool import BuiltinTool + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + + +class NovaCanvasTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Invoke AWS Bedrock Nova Canvas model for image generation + """ + # Get common parameters + prompt = tool_parameters.get("prompt", "") + image_output_s3uri = tool_parameters.get("image_output_s3uri", "").strip() + if not prompt: + return self.create_text_message("Please provide a text prompt for image generation.") + if not image_output_s3uri or urlparse(image_output_s3uri).scheme != "s3": + return self.create_text_message("Please provide an valid S3 URI for image output.") + + task_type = tool_parameters.get("task_type", "TEXT_IMAGE") + aws_region = tool_parameters.get("aws_region", "us-east-1") + + # Get common image generation config parameters + width = tool_parameters.get("width", 1024) + height = tool_parameters.get("height", 1024) + cfg_scale = tool_parameters.get("cfg_scale", 8.0) + negative_prompt = tool_parameters.get("negative_prompt", "") + seed = tool_parameters.get("seed", 0) + quality = tool_parameters.get("quality", "standard") + + # Handle S3 image if provided + image_input_s3uri = tool_parameters.get("image_input_s3uri", "") + if task_type != "TEXT_IMAGE": + if not image_input_s3uri or urlparse(image_input_s3uri).scheme != "s3": + return self.create_text_message("Please provide a valid S3 URI for image to image generation.") + + # Parse S3 URI + parsed_uri = urlparse(image_input_s3uri) + bucket = parsed_uri.netloc + key = parsed_uri.path.lstrip("/") + + # Initialize S3 client and download image + s3_client = boto3.client("s3") + response = s3_client.get_object(Bucket=bucket, Key=key) + image_data = response["Body"].read() + + # Base64 encode the image + input_image = base64.b64encode(image_data).decode("utf-8") + + try: + # Initialize Bedrock client + bedrock = boto3.client(service_name="bedrock-runtime", region_name=aws_region) + + # Base image generation config + image_generation_config = { + "width": width, + "height": height, + "cfgScale": cfg_scale, + "seed": seed, + "numberOfImages": 1, + "quality": quality, + } + + # Prepare request body based on task type + body = {"imageGenerationConfig": image_generation_config} + + if task_type == "TEXT_IMAGE": + body["taskType"] = "TEXT_IMAGE" + body["textToImageParams"] = {"text": prompt} + if negative_prompt: + body["textToImageParams"]["negativeText"] = negative_prompt + + elif task_type == "COLOR_GUIDED_GENERATION": + colors = tool_parameters.get("colors", "#ff8080-#ffb280-#ffe680-#ffe680") + if not self._validate_color_string(colors): + return self.create_text_message("Please provide valid colors in hexadecimal format.") + + body["taskType"] = "COLOR_GUIDED_GENERATION" + body["colorGuidedGenerationParams"] = { + "colors": colors.split("-"), + "referenceImage": input_image, + "text": prompt, + } + if negative_prompt: + body["colorGuidedGenerationParams"]["negativeText"] = negative_prompt + + elif task_type == "IMAGE_VARIATION": + similarity_strength = tool_parameters.get("similarity_strength", 0.5) + + body["taskType"] = "IMAGE_VARIATION" + body["imageVariationParams"] = { + "images": [input_image], + "similarityStrength": similarity_strength, + "text": prompt, + } + if negative_prompt: + body["imageVariationParams"]["negativeText"] = negative_prompt + + elif task_type == "INPAINTING": + mask_prompt = tool_parameters.get("mask_prompt") + if not mask_prompt: + return self.create_text_message("Please provide a mask prompt for image inpainting.") + + body["taskType"] = "INPAINTING" + body["inPaintingParams"] = {"image": input_image, "maskPrompt": mask_prompt, "text": prompt} + if negative_prompt: + body["inPaintingParams"]["negativeText"] = negative_prompt + + elif task_type == "OUTPAINTING": + mask_prompt = tool_parameters.get("mask_prompt") + if not mask_prompt: + return self.create_text_message("Please provide a mask prompt for image outpainting.") + outpainting_mode = tool_parameters.get("outpainting_mode", "DEFAULT") + + body["taskType"] = "OUTPAINTING" + body["outPaintingParams"] = { + "image": input_image, + "maskPrompt": mask_prompt, + "outPaintingMode": outpainting_mode, + "text": prompt, + } + if negative_prompt: + body["outPaintingParams"]["negativeText"] = negative_prompt + + elif task_type == "BACKGROUND_REMOVAL": + body["taskType"] = "BACKGROUND_REMOVAL" + body["backgroundRemovalParams"] = {"image": input_image} + + else: + return self.create_text_message(f"Unsupported task type: {task_type}") + + # Call Nova Canvas model + response = bedrock.invoke_model( + body=json.dumps(body), + modelId="amazon.nova-canvas-v1:0", + accept="application/json", + contentType="application/json", + ) + + # Process response + response_body = json.loads(response.get("body").read()) + if response_body.get("error"): + raise Exception(f"Error in model response: {response_body.get('error')}") + base64_image = response_body.get("images")[0] + + # Upload to S3 if image_output_s3uri is provided + try: + # Parse S3 URI for output + parsed_uri = urlparse(image_output_s3uri) + output_bucket = parsed_uri.netloc + output_base_path = parsed_uri.path.lstrip("/") + # Generate filename with timestamp + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + output_key = f"{output_base_path}/canvas-output-{timestamp}.png" + + # Initialize S3 client if not already done + s3_client = boto3.client("s3", region_name=aws_region) + + # Decode base64 image and upload to S3 + image_data = base64.b64decode(base64_image) + s3_client.put_object(Bucket=output_bucket, Key=output_key, Body=image_data, ContentType="image/png") + logger.info(f"Image uploaded to s3://{output_bucket}/{output_key}") + except Exception as e: + logger.exception("Failed to upload image to S3") + # Return image + return [ + self.create_text_message(f"Image is available at: s3://{output_bucket}/{output_key}"), + self.create_blob_message( + blob=base64.b64decode(base64_image), + meta={"mime_type": "image/png"}, + save_as=self.VariableKey.IMAGE.value, + ), + ] + + except Exception as e: + return self.create_text_message(f"Failed to generate image: {str(e)}") + + def _validate_color_string(self, color_string) -> bool: + color_pattern = r"^#[0-9a-fA-F]{6}(?:-#[0-9a-fA-F]{6})*$" + + if re.match(color_pattern, color_string): + return True + return False + + def get_runtime_parameters(self) -> list[ToolParameter]: + parameters = [ + ToolParameter( + name="prompt", + label=I18nObject(en_US="Prompt", zh_Hans="提示词"), + type=ToolParameter.ToolParameterType.STRING, + required=True, + form=ToolParameter.ToolParameterForm.LLM, + human_description=I18nObject( + en_US="Text description of the image you want to generate or modify", + zh_Hans="您想要生成或修改的图像的文本描述", + ), + llm_description="Describe the image you want to generate or how you want to modify the input image", + ), + ToolParameter( + name="image_input_s3uri", + label=I18nObject(en_US="Input image s3 uri", zh_Hans="输入图片的s3 uri"), + type=ToolParameter.ToolParameterType.STRING, + required=False, + form=ToolParameter.ToolParameterForm.LLM, + human_description=I18nObject(en_US="Image to be modified", zh_Hans="想要修改的图片"), + ), + ToolParameter( + name="image_output_s3uri", + label=I18nObject(en_US="Output Image S3 URI", zh_Hans="输出图片的S3 URI目录"), + type=ToolParameter.ToolParameterType.STRING, + required=True, + form=ToolParameter.ToolParameterForm.FORM, + human_description=I18nObject( + en_US="S3 URI where the generated image should be uploaded", zh_Hans="生成的图像应该上传到的S3 URI" + ), + ), + ToolParameter( + name="width", + label=I18nObject(en_US="Width", zh_Hans="宽度"), + type=ToolParameter.ToolParameterType.NUMBER, + required=False, + default=1024, + form=ToolParameter.ToolParameterForm.FORM, + human_description=I18nObject(en_US="Width of the generated image", zh_Hans="生成图像的宽度"), + ), + ToolParameter( + name="height", + label=I18nObject(en_US="Height", zh_Hans="高度"), + type=ToolParameter.ToolParameterType.NUMBER, + required=False, + default=1024, + form=ToolParameter.ToolParameterForm.FORM, + human_description=I18nObject(en_US="Height of the generated image", zh_Hans="生成图像的高度"), + ), + ToolParameter( + name="cfg_scale", + label=I18nObject(en_US="CFG Scale", zh_Hans="CFG比例"), + type=ToolParameter.ToolParameterType.NUMBER, + required=False, + default=8.0, + form=ToolParameter.ToolParameterForm.FORM, + human_description=I18nObject( + en_US="How strongly the image should conform to the prompt", zh_Hans="图像应该多大程度上符合提示词" + ), + ), + ToolParameter( + name="negative_prompt", + label=I18nObject(en_US="Negative Prompt", zh_Hans="负面提示词"), + type=ToolParameter.ToolParameterType.STRING, + required=False, + default="", + form=ToolParameter.ToolParameterForm.LLM, + human_description=I18nObject( + en_US="Things you don't want in the generated image", zh_Hans="您不想在生成的图像中出现的内容" + ), + ), + ToolParameter( + name="seed", + label=I18nObject(en_US="Seed", zh_Hans="种子值"), + type=ToolParameter.ToolParameterType.NUMBER, + required=False, + default=0, + form=ToolParameter.ToolParameterForm.FORM, + human_description=I18nObject(en_US="Random seed for image generation", zh_Hans="图像生成的随机种子"), + ), + ToolParameter( + name="aws_region", + label=I18nObject(en_US="AWS Region", zh_Hans="AWS 区域"), + type=ToolParameter.ToolParameterType.STRING, + required=False, + default="us-east-1", + form=ToolParameter.ToolParameterForm.FORM, + human_description=I18nObject(en_US="AWS region for Bedrock service", zh_Hans="Bedrock 服务的 AWS 区域"), + ), + ToolParameter( + name="task_type", + label=I18nObject(en_US="Task Type", zh_Hans="任务类型"), + type=ToolParameter.ToolParameterType.STRING, + required=False, + default="TEXT_IMAGE", + form=ToolParameter.ToolParameterForm.LLM, + human_description=I18nObject(en_US="Type of image generation task", zh_Hans="图像生成任务的类型"), + ), + ToolParameter( + name="quality", + label=I18nObject(en_US="Quality", zh_Hans="质量"), + type=ToolParameter.ToolParameterType.STRING, + required=False, + default="standard", + form=ToolParameter.ToolParameterForm.FORM, + human_description=I18nObject( + en_US="Quality of the generated image (standard or premium)", zh_Hans="生成图像的质量(标准或高级)" + ), + ), + ToolParameter( + name="colors", + label=I18nObject(en_US="Colors", zh_Hans="颜色"), + type=ToolParameter.ToolParameterType.STRING, + required=False, + form=ToolParameter.ToolParameterForm.FORM, + human_description=I18nObject( + en_US="List of colors for color-guided generation, example: #ff8080-#ffb280-#ffe680-#ffe680", + zh_Hans="颜色引导生成的颜色列表, 例子: #ff8080-#ffb280-#ffe680-#ffe680", + ), + ), + ToolParameter( + name="similarity_strength", + label=I18nObject(en_US="Similarity Strength", zh_Hans="相似度强度"), + type=ToolParameter.ToolParameterType.NUMBER, + required=False, + default=0.5, + form=ToolParameter.ToolParameterForm.FORM, + human_description=I18nObject( + en_US="How similar the generated image should be to the input image (0.0 to 1.0)", + zh_Hans="生成的图像应该与输入图像的相似程度(0.0到1.0)", + ), + ), + ToolParameter( + name="mask_prompt", + label=I18nObject(en_US="Mask Prompt", zh_Hans="蒙版提示词"), + type=ToolParameter.ToolParameterType.STRING, + required=False, + form=ToolParameter.ToolParameterForm.LLM, + human_description=I18nObject( + en_US="Text description to generate mask for inpainting/outpainting", + zh_Hans="用于生成内补绘制/外补绘制蒙版的文本描述", + ), + ), + ToolParameter( + name="outpainting_mode", + label=I18nObject(en_US="Outpainting Mode", zh_Hans="外补绘制模式"), + type=ToolParameter.ToolParameterType.STRING, + required=False, + default="DEFAULT", + form=ToolParameter.ToolParameterForm.FORM, + human_description=I18nObject( + en_US="Mode for outpainting (DEFAULT or other supported modes)", + zh_Hans="外补绘制的模式(DEFAULT或其他支持的模式)", + ), + ), + ] + + return parameters diff --git a/api/core/tools/provider/builtin/aws/tools/nova_canvas.yaml b/api/core/tools/provider/builtin/aws/tools/nova_canvas.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a72fd9c8efcce11b72f9c6812e2185546f899b6b --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/nova_canvas.yaml @@ -0,0 +1,175 @@ +identity: + name: nova_canvas + author: AWS + label: + en_US: AWS Bedrock Nova Canvas + zh_Hans: AWS Bedrock Nova Canvas + icon: icon.svg +description: + human: + en_US: A tool for generating and modifying images using AWS Bedrock's Nova Canvas model. Supports text-to-image, color-guided generation, image variation, inpainting, outpainting, and background removal. Input parameters reference https://docs.aws.amazon.com/nova/latest/userguide/image-gen-req-resp-structure.html + zh_Hans: 使用 AWS Bedrock 的 Nova Canvas 模型生成和修改图像的工具。支持文生图、颜色引导生成、图像变体、内补绘制、外补绘制和背景移除功能, 输入参数参考 https://docs.aws.amazon.com/nova/latest/userguide/image-gen-req-resp-structure.html。 + llm: Generate or modify images using AWS Bedrock's Nova Canvas model with multiple task types including text-to-image, color-guided generation, image variation, inpainting, outpainting, and background removal. +parameters: + - name: task_type + type: string + required: false + default: TEXT_IMAGE + label: + en_US: Task Type + zh_Hans: 任务类型 + human_description: + en_US: Type of image generation task (TEXT_IMAGE, COLOR_GUIDED_GENERATION, IMAGE_VARIATION, INPAINTING, OUTPAINTING, BACKGROUND_REMOVAL) + zh_Hans: 图像生成任务的类型(文生图、颜色引导生成、图像变体、内补绘制、外补绘制、背景移除) + form: llm + - name: prompt + type: string + required: true + label: + en_US: Prompt + zh_Hans: 提示词 + human_description: + en_US: Text description of the image you want to generate or modify + zh_Hans: 您想要生成或修改的图像的文本描述 + llm_description: Describe the image you want to generate or how you want to modify the input image + form: llm + - name: image_input_s3uri + type: string + required: false + label: + en_US: Input image s3 uri + zh_Hans: 输入图片的s3 uri + human_description: + en_US: The input image to modify (required for all modes except TEXT_IMAGE) + zh_Hans: 要修改的输入图像(除文生图外的所有模式都需要) + llm_description: The input image you want to modify. Required for all modes except TEXT_IMAGE. + form: llm + - name: image_output_s3uri + type: string + required: true + label: + en_US: Output S3 URI + zh_Hans: 输出S3 URI + human_description: + en_US: The S3 URI where the generated image will be saved. If provided, the image will be uploaded with name format canvas-output-{timestamp}.png + zh_Hans: 生成的图像将保存到的S3 URI。如果提供,图像将以canvas-output-{timestamp}.png的格式上传 + llm_description: Optional S3 URI where the generated image will be uploaded. The image will be saved with a timestamp-based filename. + form: form + - name: negative_prompt + type: string + required: false + label: + en_US: Negative Prompt + zh_Hans: 负面提示词 + human_description: + en_US: Things you don't want in the generated image + zh_Hans: 您不想在生成的图像中出现的内容 + form: llm + - name: width + type: number + required: false + label: + en_US: Width + zh_Hans: 宽度 + human_description: + en_US: Width of the generated image + zh_Hans: 生成图像的宽度 + form: form + default: 1024 + - name: height + type: number + required: false + label: + en_US: Height + zh_Hans: 高度 + human_description: + en_US: Height of the generated image + zh_Hans: 生成图像的高度 + form: form + default: 1024 + - name: cfg_scale + type: number + required: false + label: + en_US: CFG Scale + zh_Hans: CFG比例 + human_description: + en_US: How strongly the image should conform to the prompt + zh_Hans: 图像应该多大程度上符合提示词 + form: form + default: 8.0 + - name: seed + type: number + required: false + label: + en_US: Seed + zh_Hans: 种子值 + human_description: + en_US: Random seed for image generation + zh_Hans: 图像生成的随机种子 + form: form + default: 0 + - name: aws_region + type: string + required: false + default: us-east-1 + label: + en_US: AWS Region + zh_Hans: AWS 区域 + human_description: + en_US: AWS region for Bedrock service + zh_Hans: Bedrock 服务的 AWS 区域 + form: form + - name: quality + type: string + required: false + default: standard + label: + en_US: Quality + zh_Hans: 质量 + human_description: + en_US: Quality of the generated image (standard or premium) + zh_Hans: 生成图像的质量(标准或高级) + form: form + - name: colors + type: string + required: false + label: + en_US: Colors + zh_Hans: 颜色 + human_description: + en_US: List of colors for color-guided generation + zh_Hans: 颜色引导生成的颜色列表 + form: form + - name: similarity_strength + type: number + required: false + default: 0.5 + label: + en_US: Similarity Strength + zh_Hans: 相似度强度 + human_description: + en_US: How similar the generated image should be to the input image (0.0 to 1.0) + zh_Hans: 生成的图像应该与输入图像的相似程度(0.0到1.0) + form: form + - name: mask_prompt + type: string + required: false + label: + en_US: Mask Prompt + zh_Hans: 蒙版提示词 + human_description: + en_US: Text description to generate mask for inpainting/outpainting + zh_Hans: 用于生成内补绘制/外补绘制蒙版的文本描述 + form: llm + - name: outpainting_mode + type: string + required: false + default: DEFAULT + label: + en_US: Outpainting Mode + zh_Hans: 外补绘制模式 + human_description: + en_US: Mode for outpainting (DEFAULT or other supported modes) + zh_Hans: 外补绘制的模式(DEFAULT或其他支持的模式) + form: form diff --git a/api/core/tools/provider/builtin/aws/tools/nova_reel.py b/api/core/tools/provider/builtin/aws/tools/nova_reel.py new file mode 100644 index 0000000000000000000000000000000000000000..848df0b36bfb5a11e29344d063a610ba8efd2f6d --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/nova_reel.py @@ -0,0 +1,370 @@ +import base64 +import logging +import time +from io import BytesIO +from typing import Any, Optional, Union +from urllib.parse import urlparse + +import boto3 +from botocore.exceptions import ClientError +from PIL import Image + +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter +from core.tools.tool.builtin_tool import BuiltinTool + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +NOVA_REEL_DEFAULT_REGION = "us-east-1" +NOVA_REEL_DEFAULT_DIMENSION = "1280x720" +NOVA_REEL_DEFAULT_FPS = 24 +NOVA_REEL_DEFAULT_DURATION = 6 +NOVA_REEL_MODEL_ID = "amazon.nova-reel-v1:0" +NOVA_REEL_STATUS_CHECK_INTERVAL = 5 + +# Image requirements +NOVA_REEL_REQUIRED_IMAGE_WIDTH = 1280 +NOVA_REEL_REQUIRED_IMAGE_HEIGHT = 720 +NOVA_REEL_REQUIRED_IMAGE_MODE = "RGB" + + +class NovaReelTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Invoke AWS Bedrock Nova Reel model for video generation. + + Args: + user_id: The ID of the user making the request + tool_parameters: Dictionary containing the tool parameters + + Returns: + ToolInvokeMessage containing either the video content or status information + """ + try: + # Validate and extract parameters + params = self._validate_and_extract_parameters(tool_parameters) + if isinstance(params, ToolInvokeMessage): + return params + + # Initialize AWS clients + bedrock, s3_client = self._initialize_aws_clients(params["aws_region"]) + + # Prepare model input + model_input = self._prepare_model_input(params, s3_client) + if isinstance(model_input, ToolInvokeMessage): + return model_input + + # Start video generation + invocation = self._start_video_generation(bedrock, model_input, params["video_output_s3uri"]) + invocation_arn = invocation["invocationArn"] + + # Handle async/sync mode + return self._handle_generation_mode(bedrock, s3_client, invocation_arn, params["async_mode"]) + + except ClientError as e: + error_code = e.response.get("Error", {}).get("Code", "Unknown") + error_message = e.response.get("Error", {}).get("Message", str(e)) + logger.exception(f"AWS API error: {error_code} - {error_message}") + return self.create_text_message(f"AWS service error: {error_code} - {error_message}") + except Exception as e: + logger.error(f"Unexpected error in video generation: {str(e)}", exc_info=True) + return self.create_text_message(f"Failed to generate video: {str(e)}") + + def _validate_and_extract_parameters( + self, tool_parameters: dict[str, Any] + ) -> Union[dict[str, Any], ToolInvokeMessage]: + """Validate and extract parameters from the input dictionary.""" + prompt = tool_parameters.get("prompt", "") + video_output_s3uri = tool_parameters.get("video_output_s3uri", "").strip() + + # Validate required parameters + if not prompt: + return self.create_text_message("Please provide a text prompt for video generation.") + if not video_output_s3uri: + return self.create_text_message("Please provide an S3 URI for video output.") + + # Validate S3 URI format + if not video_output_s3uri.startswith("s3://"): + return self.create_text_message("Invalid S3 URI format. Must start with 's3://'") + + # Ensure S3 URI ends with '/' + video_output_s3uri = video_output_s3uri if video_output_s3uri.endswith("/") else video_output_s3uri + "/" + + return { + "prompt": prompt, + "video_output_s3uri": video_output_s3uri, + "image_input_s3uri": tool_parameters.get("image_input_s3uri", "").strip(), + "aws_region": tool_parameters.get("aws_region", NOVA_REEL_DEFAULT_REGION), + "dimension": tool_parameters.get("dimension", NOVA_REEL_DEFAULT_DIMENSION), + "seed": int(tool_parameters.get("seed", 0)), + "fps": int(tool_parameters.get("fps", NOVA_REEL_DEFAULT_FPS)), + "duration": int(tool_parameters.get("duration", NOVA_REEL_DEFAULT_DURATION)), + "async_mode": bool(tool_parameters.get("async", True)), + } + + def _initialize_aws_clients(self, region: str) -> tuple[Any, Any]: + """Initialize AWS Bedrock and S3 clients.""" + bedrock = boto3.client(service_name="bedrock-runtime", region_name=region) + s3_client = boto3.client("s3", region_name=region) + return bedrock, s3_client + + def _prepare_model_input(self, params: dict[str, Any], s3_client: Any) -> Union[dict[str, Any], ToolInvokeMessage]: + """Prepare the input for the Nova Reel model.""" + model_input = { + "taskType": "TEXT_VIDEO", + "textToVideoParams": {"text": params["prompt"]}, + "videoGenerationConfig": { + "durationSeconds": params["duration"], + "fps": params["fps"], + "dimension": params["dimension"], + "seed": params["seed"], + }, + } + + # Add image if provided + if params["image_input_s3uri"]: + try: + image_data = self._get_image_from_s3(s3_client, params["image_input_s3uri"]) + if not image_data: + return self.create_text_message("Failed to retrieve image from S3") + + # Process and validate image + processed_image = self._process_and_validate_image(image_data) + if isinstance(processed_image, ToolInvokeMessage): + return processed_image + + # Convert processed image to base64 + img_buffer = BytesIO() + processed_image.save(img_buffer, format="PNG") + img_buffer.seek(0) + input_image_base64 = base64.b64encode(img_buffer.getvalue()).decode("utf-8") + + model_input["textToVideoParams"]["images"] = [ + {"format": "png", "source": {"bytes": input_image_base64}} + ] + except Exception as e: + logger.error(f"Error processing input image: {str(e)}", exc_info=True) + return self.create_text_message(f"Failed to process input image: {str(e)}") + + return model_input + + def _process_and_validate_image(self, image_data: bytes) -> Union[Image.Image, ToolInvokeMessage]: + """ + Process and validate the input image according to Nova Reel requirements. + + Requirements: + - Must be 1280x720 pixels + - Must be RGB format (8 bits per channel) + - If PNG, alpha channel must not have transparent/translucent pixels + """ + try: + # Open image + img = Image.open(BytesIO(image_data)) + + # Convert RGBA to RGB if needed, ensuring no transparency + if img.mode == "RGBA": + # Check for transparency + if img.getchannel("A").getextrema()[0] < 255: + return self.create_text_message( + "PNG image contains transparent or translucent pixels, which is not supported. " + "Please provide an image without transparency." + ) + # Convert to RGB + img = img.convert("RGB") + elif img.mode != "RGB": + # Convert any other mode to RGB + img = img.convert("RGB") + + # Validate/adjust dimensions + if img.size != (NOVA_REEL_REQUIRED_IMAGE_WIDTH, NOVA_REEL_REQUIRED_IMAGE_HEIGHT): + logger.warning( + f"Image dimensions {img.size} do not match required dimensions " + f"({NOVA_REEL_REQUIRED_IMAGE_WIDTH}x{NOVA_REEL_REQUIRED_IMAGE_HEIGHT}). Resizing..." + ) + img = img.resize( + (NOVA_REEL_REQUIRED_IMAGE_WIDTH, NOVA_REEL_REQUIRED_IMAGE_HEIGHT), Image.Resampling.LANCZOS + ) + + # Validate bit depth + if img.mode != NOVA_REEL_REQUIRED_IMAGE_MODE: + return self.create_text_message( + f"Image must be in {NOVA_REEL_REQUIRED_IMAGE_MODE} mode with 8 bits per channel" + ) + + return img + + except Exception as e: + logger.error(f"Error processing image: {str(e)}", exc_info=True) + return self.create_text_message( + "Failed to process image. Please ensure the image is a valid JPEG or PNG file." + ) + + def _get_image_from_s3(self, s3_client: Any, s3_uri: str) -> Optional[bytes]: + """Download and return image data from S3.""" + parsed_uri = urlparse(s3_uri) + bucket = parsed_uri.netloc + key = parsed_uri.path.lstrip("/") + + response = s3_client.get_object(Bucket=bucket, Key=key) + return response["Body"].read() + + def _start_video_generation(self, bedrock: Any, model_input: dict[str, Any], output_s3uri: str) -> dict[str, Any]: + """Start the async video generation process.""" + return bedrock.start_async_invoke( + modelId=NOVA_REEL_MODEL_ID, + modelInput=model_input, + outputDataConfig={"s3OutputDataConfig": {"s3Uri": output_s3uri}}, + ) + + def _handle_generation_mode( + self, bedrock: Any, s3_client: Any, invocation_arn: str, async_mode: bool + ) -> ToolInvokeMessage: + """Handle async or sync video generation mode.""" + invocation_response = bedrock.get_async_invoke(invocationArn=invocation_arn) + video_path = invocation_response["outputDataConfig"]["s3OutputDataConfig"]["s3Uri"] + video_uri = f"{video_path}/output.mp4" + + if async_mode: + return self.create_text_message( + f"Video generation started.\nInvocation ARN: {invocation_arn}\nVideo will be available at: {video_uri}" + ) + + return self._wait_for_completion(bedrock, s3_client, invocation_arn) + + def _wait_for_completion(self, bedrock: Any, s3_client: Any, invocation_arn: str) -> ToolInvokeMessage: + """Wait for video generation completion and handle the result.""" + while True: + status_response = bedrock.get_async_invoke(invocationArn=invocation_arn) + status = status_response["status"] + video_path = status_response["outputDataConfig"]["s3OutputDataConfig"]["s3Uri"] + + if status == "Completed": + return self._handle_completed_video(s3_client, video_path) + elif status == "Failed": + failure_message = status_response.get("failureMessage", "Unknown error") + return self.create_text_message(f"Video generation failed.\nError: {failure_message}") + elif status == "InProgress": + time.sleep(NOVA_REEL_STATUS_CHECK_INTERVAL) + else: + return self.create_text_message(f"Unexpected status: {status}") + + def _handle_completed_video(self, s3_client: Any, video_path: str) -> ToolInvokeMessage: + """Handle completed video generation and return the result.""" + parsed_uri = urlparse(video_path) + bucket = parsed_uri.netloc + key = parsed_uri.path.lstrip("/") + "/output.mp4" + + try: + response = s3_client.get_object(Bucket=bucket, Key=key) + video_content = response["Body"].read() + return [ + self.create_text_message(f"Video is available at: {video_path}/output.mp4"), + self.create_blob_message(blob=video_content, meta={"mime_type": "video/mp4"}, save_as="output.mp4"), + ] + except Exception as e: + logger.error(f"Error downloading video: {str(e)}", exc_info=True) + return self.create_text_message( + f"Video generation completed but failed to download video: {str(e)}\n" + f"Video is available at: s3://{bucket}/{key}" + ) + + def get_runtime_parameters(self) -> list[ToolParameter]: + """Define the tool's runtime parameters.""" + parameters = [ + ToolParameter( + name="prompt", + label=I18nObject(en_US="Prompt", zh_Hans="提示词"), + type=ToolParameter.ToolParameterType.STRING, + required=True, + form=ToolParameter.ToolParameterForm.LLM, + human_description=I18nObject( + en_US="Text description of the video you want to generate", zh_Hans="您想要生成的视频的文本描述" + ), + llm_description="Describe the video you want to generate", + ), + ToolParameter( + name="video_output_s3uri", + label=I18nObject(en_US="Output S3 URI", zh_Hans="输出S3 URI"), + type=ToolParameter.ToolParameterType.STRING, + required=True, + form=ToolParameter.ToolParameterForm.FORM, + human_description=I18nObject( + en_US="S3 URI where the generated video will be stored", zh_Hans="生成的视频将存储的S3 URI" + ), + ), + ToolParameter( + name="dimension", + label=I18nObject(en_US="Dimension", zh_Hans="尺寸"), + type=ToolParameter.ToolParameterType.STRING, + required=False, + default=NOVA_REEL_DEFAULT_DIMENSION, + form=ToolParameter.ToolParameterForm.FORM, + human_description=I18nObject(en_US="Video dimensions (width x height)", zh_Hans="视频尺寸(宽 x 高)"), + ), + ToolParameter( + name="duration", + label=I18nObject(en_US="Duration", zh_Hans="时长"), + type=ToolParameter.ToolParameterType.NUMBER, + required=False, + default=NOVA_REEL_DEFAULT_DURATION, + form=ToolParameter.ToolParameterForm.FORM, + human_description=I18nObject(en_US="Video duration in seconds", zh_Hans="视频时长(秒)"), + ), + ToolParameter( + name="seed", + label=I18nObject(en_US="Seed", zh_Hans="种子值"), + type=ToolParameter.ToolParameterType.NUMBER, + required=False, + default=0, + form=ToolParameter.ToolParameterForm.FORM, + human_description=I18nObject(en_US="Random seed for video generation", zh_Hans="视频生成的随机种子"), + ), + ToolParameter( + name="fps", + label=I18nObject(en_US="FPS", zh_Hans="帧率"), + type=ToolParameter.ToolParameterType.NUMBER, + required=False, + default=NOVA_REEL_DEFAULT_FPS, + form=ToolParameter.ToolParameterForm.FORM, + human_description=I18nObject( + en_US="Frames per second for the generated video", zh_Hans="生成视频的每秒帧数" + ), + ), + ToolParameter( + name="async", + label=I18nObject(en_US="Async Mode", zh_Hans="异步模式"), + type=ToolParameter.ToolParameterType.BOOLEAN, + required=False, + default=True, + form=ToolParameter.ToolParameterForm.LLM, + human_description=I18nObject( + en_US="Whether to run in async mode (return immediately) or sync mode (wait for completion)", + zh_Hans="是否以异步模式运行(立即返回)或同步模式(等待完成)", + ), + ), + ToolParameter( + name="aws_region", + label=I18nObject(en_US="AWS Region", zh_Hans="AWS 区域"), + type=ToolParameter.ToolParameterType.STRING, + required=False, + default=NOVA_REEL_DEFAULT_REGION, + form=ToolParameter.ToolParameterForm.FORM, + human_description=I18nObject(en_US="AWS region for Bedrock service", zh_Hans="Bedrock 服务的 AWS 区域"), + ), + ToolParameter( + name="image_input_s3uri", + label=I18nObject(en_US="Input Image S3 URI", zh_Hans="输入图像S3 URI"), + type=ToolParameter.ToolParameterType.STRING, + required=False, + form=ToolParameter.ToolParameterForm.LLM, + human_description=I18nObject( + en_US="S3 URI of the input image (1280x720 JPEG/PNG) to use as first frame", + zh_Hans="用作第一帧的输入图像(1280x720 JPEG/PNG)的S3 URI", + ), + ), + ] + + return parameters diff --git a/api/core/tools/provider/builtin/aws/tools/nova_reel.yaml b/api/core/tools/provider/builtin/aws/tools/nova_reel.yaml new file mode 100644 index 0000000000000000000000000000000000000000..16df5ba5c9d1e39e9367f7f9d4e21881a8f30189 --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/nova_reel.yaml @@ -0,0 +1,124 @@ +identity: + name: nova_reel + author: AWS + label: + en_US: AWS Bedrock Nova Reel + zh_Hans: AWS Bedrock Nova Reel + icon: icon.svg +description: + human: + en_US: A tool for generating videos using AWS Bedrock's Nova Reel model. Supports text-to-video generation and image-to-video generation with customizable parameters like duration, FPS, and dimensions. Input parameters reference https://docs.aws.amazon.com/nova/latest/userguide/video-generation.html + zh_Hans: 使用 AWS Bedrock 的 Nova Reel 模型生成视频的工具。支持文本生成视频和图像生成视频功能,可自定义持续时间、帧率和尺寸等参数。输入参数参考 https://docs.aws.amazon.com/nova/latest/userguide/video-generation.html + llm: Generate videos using AWS Bedrock's Nova Reel model with support for both text-to-video and image-to-video generation, allowing customization of video properties like duration, frame rate, and resolution. + +parameters: + - name: prompt + type: string + required: true + label: + en_US: Prompt + zh_Hans: 提示词 + human_description: + en_US: Text description of the video you want to generate + zh_Hans: 您想要生成的视频的文本描述 + llm_description: Describe the video you want to generate + form: llm + + - name: video_output_s3uri + type: string + required: true + label: + en_US: Output S3 URI + zh_Hans: 输出S3 URI + human_description: + en_US: S3 URI where the generated video will be stored + zh_Hans: 生成的视频将存储的S3 URI + form: form + + - name: dimension + type: string + required: false + default: 1280x720 + label: + en_US: Dimension + zh_Hans: 尺寸 + human_description: + en_US: Video dimensions (width x height) + zh_Hans: 视频尺寸(宽 x 高) + form: form + + - name: duration + type: number + required: false + default: 6 + label: + en_US: Duration + zh_Hans: 时长 + human_description: + en_US: Video duration in seconds + zh_Hans: 视频时长(秒) + form: form + + - name: seed + type: number + required: false + default: 0 + label: + en_US: Seed + zh_Hans: 种子值 + human_description: + en_US: Random seed for video generation + zh_Hans: 视频生成的随机种子 + form: form + + - name: fps + type: number + required: false + default: 24 + label: + en_US: FPS + zh_Hans: 帧率 + human_description: + en_US: Frames per second for the generated video + zh_Hans: 生成视频的每秒帧数 + form: form + + - name: async + type: boolean + required: false + default: true + label: + en_US: Async Mode + zh_Hans: 异步模式 + human_description: + en_US: Whether to run in async mode (return immediately) or sync mode (wait for completion) + zh_Hans: 是否以异步模式运行(立即返回)或同步模式(等待完成) + form: llm + + - name: aws_region + type: string + required: false + default: us-east-1 + label: + en_US: AWS Region + zh_Hans: AWS 区域 + human_description: + en_US: AWS region for Bedrock service + zh_Hans: Bedrock 服务的 AWS 区域 + form: form + + - name: image_input_s3uri + type: string + required: false + label: + en_US: Input Image S3 URI + zh_Hans: 输入图像S3 URI + human_description: + en_US: S3 URI of the input image (1280x720 JPEG/PNG) to use as first frame + zh_Hans: 用作第一帧的输入图像(1280x720 JPEG/PNG)的S3 URI + form: llm + +development: + dependencies: + - boto3 + - pillow diff --git a/api/core/tools/provider/builtin/aws/tools/s3_operator.py b/api/core/tools/provider/builtin/aws/tools/s3_operator.py new file mode 100644 index 0000000000000000000000000000000000000000..e4026b07a873106e9fa5813f817df7701bf6e802 --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/s3_operator.py @@ -0,0 +1,80 @@ +from typing import Any, Union +from urllib.parse import urlparse + +import boto3 + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class S3Operator(BuiltinTool): + s3_client: Any = None + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + try: + # Initialize S3 client if not already done + if not self.s3_client: + aws_region = tool_parameters.get("aws_region") + if aws_region: + self.s3_client = boto3.client("s3", region_name=aws_region) + else: + self.s3_client = boto3.client("s3") + + # Parse S3 URI + s3_uri = tool_parameters.get("s3_uri") + if not s3_uri: + return self.create_text_message("s3_uri parameter is required") + + parsed_uri = urlparse(s3_uri) + if parsed_uri.scheme != "s3": + return self.create_text_message("Invalid S3 URI format. Must start with 's3://'") + + bucket = parsed_uri.netloc + # Remove leading slash from key + key = parsed_uri.path.lstrip("/") + + operation_type = tool_parameters.get("operation_type", "read") + generate_presign_url = tool_parameters.get("generate_presign_url", False) + presign_expiry = int(tool_parameters.get("presign_expiry", 3600)) # default 1 hour + + if operation_type == "write": + text_content = tool_parameters.get("text_content") + if not text_content: + return self.create_text_message("text_content parameter is required for write operation") + + # Write content to S3 + self.s3_client.put_object(Bucket=bucket, Key=key, Body=text_content.encode("utf-8")) + result = f"s3://{bucket}/{key}" + + # Generate presigned URL for the written object if requested + if generate_presign_url: + result = self.s3_client.generate_presigned_url( + "get_object", Params={"Bucket": bucket, "Key": key}, ExpiresIn=presign_expiry + ) + + else: # read operation + # Get object from S3 + response = self.s3_client.get_object(Bucket=bucket, Key=key) + result = response["Body"].read().decode("utf-8") + + # Generate presigned URL if requested + if generate_presign_url: + result = self.s3_client.generate_presigned_url( + "get_object", Params={"Bucket": bucket, "Key": key}, ExpiresIn=presign_expiry + ) + + return self.create_text_message(text=result) + + except self.s3_client.exceptions.NoSuchBucket: + return self.create_text_message(f"Bucket '{bucket}' does not exist") + except self.s3_client.exceptions.NoSuchKey: + return self.create_text_message(f"Object '{key}' does not exist in bucket '{bucket}'") + except Exception as e: + return self.create_text_message(f"Exception: {str(e)}") diff --git a/api/core/tools/provider/builtin/aws/tools/s3_operator.yaml b/api/core/tools/provider/builtin/aws/tools/s3_operator.yaml new file mode 100644 index 0000000000000000000000000000000000000000..642fc2966e9b6d908485fa83cdcd9a00e2250ce8 --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/s3_operator.yaml @@ -0,0 +1,98 @@ +identity: + name: s3_operator + author: AWS + label: + en_US: AWS S3 Operator + zh_Hans: AWS S3 读写器 + pt_BR: AWS S3 Operator + icon: icon.svg +description: + human: + en_US: AWS S3 Writer and Reader + zh_Hans: 读写S3 bucket中的文件 + pt_BR: AWS S3 Writer and Reader + llm: AWS S3 Writer and Reader +parameters: + - name: text_content + type: string + required: false + label: + en_US: The text to write + zh_Hans: 待写入的文本 + pt_BR: The text to write + human_description: + en_US: The text to write + zh_Hans: 待写入的文本 + pt_BR: The text to write + llm_description: The text to write + form: llm + - name: s3_uri + type: string + required: true + label: + en_US: s3 uri + zh_Hans: s3 uri + pt_BR: s3 uri + human_description: + en_US: s3 uri + zh_Hans: s3 uri + pt_BR: s3 uri + llm_description: s3 uri + form: llm + - name: aws_region + type: string + required: true + label: + en_US: region of bucket + zh_Hans: bucket 所在的region + pt_BR: region of bucket + human_description: + en_US: region of bucket + zh_Hans: bucket 所在的region + pt_BR: region of bucket + llm_description: region of bucket + form: form + - name: operation_type + type: select + required: true + label: + en_US: operation type + zh_Hans: 操作类型 + pt_BR: operation type + human_description: + en_US: operation type + zh_Hans: 操作类型 + pt_BR: operation type + default: read + options: + - value: read + label: + en_US: read + zh_Hans: 读 + - value: write + label: + en_US: write + zh_Hans: 写 + form: form + - name: generate_presign_url + type: boolean + required: false + label: + en_US: Generate presigned URL + zh_Hans: 生成预签名URL + human_description: + en_US: Whether to generate a presigned URL for the S3 object + zh_Hans: 是否生成S3对象的预签名URL + default: false + form: form + - name: presign_expiry + type: number + required: false + label: + en_US: Presigned URL expiration time + zh_Hans: 预签名URL有效期 + human_description: + en_US: Expiration time in seconds for the presigned URL + zh_Hans: 预签名URL的有效期(秒) + default: 3600 + form: form diff --git a/api/core/tools/provider/builtin/aws/tools/sagemaker_chinese_toxicity_detector.py b/api/core/tools/provider/builtin/aws/tools/sagemaker_chinese_toxicity_detector.py new file mode 100644 index 0000000000000000000000000000000000000000..3d88f28dbd2fc72b8f7beb959406e89aebfca7d2 --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/sagemaker_chinese_toxicity_detector.py @@ -0,0 +1,67 @@ +import json +from typing import Any, Union + +import boto3 + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +# Define label mappings +LABEL_MAPPING = {0: "SAFE", 1: "NO_SAFE"} + + +class ContentModerationTool(BuiltinTool): + sagemaker_client: Any = None + sagemaker_endpoint: str = None + + def _invoke_sagemaker(self, payload: dict, endpoint: str): + response = self.sagemaker_client.invoke_endpoint( + EndpointName=endpoint, + Body=json.dumps(payload), + ContentType="application/json", + ) + # Parse response + response_body = response["Body"].read().decode("utf8") + + json_obj = json.loads(response_body) + + # Handle nested JSON if present + if isinstance(json_obj, dict) and "body" in json_obj: + body_content = json.loads(json_obj["body"]) + prediction_result = body_content.get("prediction") + else: + prediction_result = json_obj.get("prediction") + + # Map labels and return + result = LABEL_MAPPING.get(prediction_result, "NO_SAFE") # If not found in mapping, default to NO_SAFE + return result + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + try: + if not self.sagemaker_client: + aws_region = tool_parameters.get("aws_region") + if aws_region: + self.sagemaker_client = boto3.client("sagemaker-runtime", region_name=aws_region) + else: + self.sagemaker_client = boto3.client("sagemaker-runtime") + + if not self.sagemaker_endpoint: + self.sagemaker_endpoint = tool_parameters.get("sagemaker_endpoint") + + content_text = tool_parameters.get("content_text") + + payload = {"text": content_text} + + result = self._invoke_sagemaker(payload, self.sagemaker_endpoint) + + return self.create_text_message(text=result) + + except Exception as e: + return self.create_text_message(f"Exception {str(e)}") diff --git a/api/core/tools/provider/builtin/aws/tools/sagemaker_chinese_toxicity_detector.yaml b/api/core/tools/provider/builtin/aws/tools/sagemaker_chinese_toxicity_detector.yaml new file mode 100644 index 0000000000000000000000000000000000000000..76dcb89632f2707ac13f3d9fd9c339d3678b01b1 --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/sagemaker_chinese_toxicity_detector.yaml @@ -0,0 +1,46 @@ +identity: + name: chinese_toxicity_detector + author: AWS + label: + en_US: Chinese Toxicity Detector + zh_Hans: 中文有害内容检测 + icon: icon.svg +description: + human: + en_US: A tool to detect Chinese toxicity + zh_Hans: 检测中文有害内容的工具 + llm: A tool that checks if Chinese content is safe for work +parameters: + - name: sagemaker_endpoint + type: string + required: true + label: + en_US: sagemaker endpoint for moderation + zh_Hans: 内容审核的SageMaker端点 + human_description: + en_US: sagemaker endpoint for content moderation + zh_Hans: 内容审核的SageMaker端点 + llm_description: sagemaker endpoint for content moderation + form: form + - name: content_text + type: string + required: true + label: + en_US: content text + zh_Hans: 待审核文本 + human_description: + en_US: text content to be moderated + zh_Hans: 需要审核的文本内容 + llm_description: text content to be moderated + form: llm + - name: aws_region + type: string + required: false + label: + en_US: region of sagemaker endpoint + zh_Hans: SageMaker 端点所在的region + human_description: + en_US: region of sagemaker endpoint + zh_Hans: SageMaker 端点所在的region + llm_description: region of sagemaker endpoint + form: form diff --git a/api/core/tools/provider/builtin/aws/tools/sagemaker_text_rerank.py b/api/core/tools/provider/builtin/aws/tools/sagemaker_text_rerank.py new file mode 100644 index 0000000000000000000000000000000000000000..8320bd84efa44007e92f51574293b558f398d649 --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/sagemaker_text_rerank.py @@ -0,0 +1,79 @@ +import json +import operator +from typing import Any, Union + +import boto3 # type: ignore + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class SageMakerReRankTool(BuiltinTool): + sagemaker_client: Any = None + sagemaker_endpoint: str = None + + def _sagemaker_rerank(self, query_input: str, docs: list[str], rerank_endpoint: str): + inputs = [query_input] * len(docs) + response_model = self.sagemaker_client.invoke_endpoint( + EndpointName=rerank_endpoint, + Body=json.dumps({"inputs": inputs, "docs": docs}), + ContentType="application/json", + ) + json_str = response_model["Body"].read().decode("utf8") + json_obj = json.loads(json_str) + scores = json_obj["scores"] + return scores if isinstance(scores, list) else [scores] + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + line = 0 + try: + if not self.sagemaker_client: + aws_region = tool_parameters.get("aws_region") + if aws_region: + self.sagemaker_client = boto3.client("sagemaker-runtime", region_name=aws_region) + else: + self.sagemaker_client = boto3.client("sagemaker-runtime") + + line = 1 + if not self.sagemaker_endpoint: + self.sagemaker_endpoint = tool_parameters.get("sagemaker_endpoint") + + line = 2 + topk = tool_parameters.get("topk", 5) + + line = 3 + query = tool_parameters.get("query", "") + if not query: + return self.create_text_message("Please input query") + + line = 4 + candidate_texts = tool_parameters.get("candidate_texts") + if not candidate_texts: + return self.create_text_message("Please input candidate_texts") + + line = 5 + candidate_docs = json.loads(candidate_texts) + docs = [item.get("content") for item in candidate_docs] + + line = 6 + scores = self._sagemaker_rerank(query_input=query, docs=docs, rerank_endpoint=self.sagemaker_endpoint) + + line = 7 + for idx in range(len(candidate_docs)): + candidate_docs[idx]["score"] = scores[idx] + + line = 8 + sorted_candidate_docs = sorted(candidate_docs, key=operator.itemgetter("score"), reverse=True) + + line = 9 + return [self.create_json_message(res) for res in sorted_candidate_docs[:topk]] + + except Exception as e: + return self.create_text_message(f"Exception {str(e)}, line : {line}") diff --git a/api/core/tools/provider/builtin/aws/tools/sagemaker_text_rerank.yaml b/api/core/tools/provider/builtin/aws/tools/sagemaker_text_rerank.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d1dfdb9f84a858062a9722f5ecc9cc3c43118c8c --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/sagemaker_text_rerank.yaml @@ -0,0 +1,82 @@ +identity: + name: sagemaker_text_rerank + author: AWS + label: + en_US: SagemakerRerank + zh_Hans: Sagemaker重排序 + pt_BR: SagemakerRerank + icon: icon.svg +description: + human: + en_US: A tool for performing text similarity ranking. You can find deploy notebook on Github Repo - https://github.com/aws-samples/dify-aws-tool + zh_Hans: Sagemaker重排序工具, 请参考 Github Repo - https://github.com/aws-samples/dify-aws-tool上的部署脚本 + pt_BR: A tool for performing text similarity ranking. + llm: A tool for performing text similarity ranking. You can find deploy notebook on Github Repo - https://github.com/aws-samples/dify-aws-tool +parameters: + - name: sagemaker_endpoint + type: string + required: true + label: + en_US: sagemaker endpoint for reranking + zh_Hans: 重排序的SageMaker 端点 + pt_BR: sagemaker endpoint for reranking + human_description: + en_US: sagemaker endpoint for reranking + zh_Hans: 重排序的SageMaker 端点 + pt_BR: sagemaker endpoint for reranking + llm_description: sagemaker endpoint for reranking + form: form + - name: query + type: string + required: true + label: + en_US: Query string + zh_Hans: 查询语句 + pt_BR: Query string + human_description: + en_US: key words for searching + zh_Hans: 查询关键词 + pt_BR: key words for searching + llm_description: key words for searching + form: llm + - name: candidate_texts + type: string + required: true + label: + en_US: text candidates + zh_Hans: 候选文本 + pt_BR: text candidates + human_description: + en_US: searched candidates by query + zh_Hans: 查询文本搜到候选文本 + pt_BR: searched candidates by query + llm_description: searched candidates by query + form: llm + - name: topk + type: number + required: false + form: form + label: + en_US: Limit for results count + zh_Hans: 返回个数限制 + pt_BR: Limit for results count + human_description: + en_US: Limit for results count + zh_Hans: 返回个数限制 + pt_BR: Limit for results count + min: 1 + max: 10 + default: 5 + - name: aws_region + type: string + required: false + label: + en_US: region of sagemaker endpoint + zh_Hans: SageMaker 端点所在的region + pt_BR: region of sagemaker endpoint + human_description: + en_US: region of sagemaker endpoint + zh_Hans: SageMaker 端点所在的region + pt_BR: region of sagemaker endpoint + llm_description: region of sagemaker endpoint + form: form diff --git a/api/core/tools/provider/builtin/aws/tools/sagemaker_tts.py b/api/core/tools/provider/builtin/aws/tools/sagemaker_tts.py new file mode 100644 index 0000000000000000000000000000000000000000..55cff89798a4eb6c567e45a68062e94cda8c4be9 --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/sagemaker_tts.py @@ -0,0 +1,101 @@ +import json +from enum import Enum +from typing import Any, Optional, Union + +import boto3 # type: ignore + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class TTSModelType(Enum): + PresetVoice = "PresetVoice" + CloneVoice = "CloneVoice" + CloneVoice_CrossLingual = "CloneVoice_CrossLingual" + InstructVoice = "InstructVoice" + + +class SageMakerTTSTool(BuiltinTool): + sagemaker_client: Any = None + sagemaker_endpoint: str | None = None + s3_client: Any = None + comprehend_client: Any = None + + def _detect_lang_code(self, content: str, map_dict: Optional[dict] = None): + map_dict = {"zh": "<|zh|>", "en": "<|en|>", "ja": "<|jp|>", "zh-TW": "<|yue|>", "ko": "<|ko|>"} + + response = self.comprehend_client.detect_dominant_language(Text=content) + language_code = response["Languages"][0]["LanguageCode"] + return map_dict.get(language_code, "<|zh|>") + + def _build_tts_payload( + self, + model_type: str, + content_text: str, + model_role: str, + prompt_text: str, + prompt_audio: str, + instruct_text: str, + ): + if model_type == TTSModelType.PresetVoice.value and model_role: + return {"tts_text": content_text, "role": model_role} + if model_type == TTSModelType.CloneVoice.value and prompt_text and prompt_audio: + return {"tts_text": content_text, "prompt_text": prompt_text, "prompt_audio": prompt_audio} + if model_type == TTSModelType.CloneVoice_CrossLingual.value and prompt_audio: + lang_tag = self._detect_lang_code(content_text) + return {"tts_text": f"{content_text}", "prompt_audio": prompt_audio, "lang_tag": lang_tag} + if model_type == TTSModelType.InstructVoice.value and instruct_text and model_role: + return {"tts_text": content_text, "role": model_role, "instruct_text": instruct_text} + + raise RuntimeError(f"Invalid params for {model_type}") + + def _invoke_sagemaker(self, payload: dict, endpoint: str): + response_model = self.sagemaker_client.invoke_endpoint( + EndpointName=endpoint, + Body=json.dumps(payload), + ContentType="application/json", + ) + json_str = response_model["Body"].read().decode("utf8") + json_obj = json.loads(json_str) + return json_obj + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + try: + if not self.sagemaker_client: + aws_region = tool_parameters.get("aws_region") + if aws_region: + self.sagemaker_client = boto3.client("sagemaker-runtime", region_name=aws_region) + self.s3_client = boto3.client("s3", region_name=aws_region) + self.comprehend_client = boto3.client("comprehend", region_name=aws_region) + else: + self.sagemaker_client = boto3.client("sagemaker-runtime") + self.s3_client = boto3.client("s3") + self.comprehend_client = boto3.client("comprehend") + + if not self.sagemaker_endpoint: + self.sagemaker_endpoint = tool_parameters.get("sagemaker_endpoint") + + tts_text = tool_parameters.get("tts_text") + tts_infer_type = tool_parameters.get("tts_infer_type") + + voice = tool_parameters.get("voice") + mock_voice_audio = tool_parameters.get("mock_voice_audio") + mock_voice_text = tool_parameters.get("mock_voice_text") + voice_instruct_prompt = tool_parameters.get("voice_instruct_prompt") + payload = self._build_tts_payload( + tts_infer_type, tts_text, voice, mock_voice_text, mock_voice_audio, voice_instruct_prompt + ) + + result = self._invoke_sagemaker(payload, self.sagemaker_endpoint) + + return self.create_text_message(text=result["s3_presign_url"]) + + except Exception as e: + return self.create_text_message(f"Exception {str(e)}") diff --git a/api/core/tools/provider/builtin/aws/tools/sagemaker_tts.yaml b/api/core/tools/provider/builtin/aws/tools/sagemaker_tts.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a6a61dd4aa519a24bf21cb9cae17d34f6889fbbd --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/sagemaker_tts.yaml @@ -0,0 +1,149 @@ +identity: + name: sagemaker_tts + author: AWS + label: + en_US: SagemakerTTS + zh_Hans: Sagemaker语音合成 + pt_BR: SagemakerTTS + icon: icon.svg +description: + human: + en_US: A tool for Speech synthesis - https://github.com/aws-samples/dify-aws-tool + zh_Hans: Sagemaker语音合成工具, 请参考 Github Repo - https://github.com/aws-samples/dify-aws-tool上的部署脚本 + pt_BR: A tool for Speech synthesis. + llm: A tool for Speech synthesis. You can find deploy notebook on Github Repo - https://github.com/aws-samples/dify-aws-tool +parameters: + - name: sagemaker_endpoint + type: string + required: true + label: + en_US: sagemaker endpoint for tts + zh_Hans: 语音生成的SageMaker端点 + pt_BR: sagemaker endpoint for tts + human_description: + en_US: sagemaker endpoint for tts + zh_Hans: 语音生成的SageMaker端点 + pt_BR: sagemaker endpoint for tts + llm_description: sagemaker endpoint for tts + form: form + - name: tts_text + type: string + required: true + label: + en_US: tts text + zh_Hans: 语音合成原文 + pt_BR: tts text + human_description: + en_US: tts text + zh_Hans: 语音合成原文 + pt_BR: tts text + llm_description: tts text + form: llm + - name: tts_infer_type + type: select + required: false + label: + en_US: tts infer type + zh_Hans: 合成方式 + pt_BR: tts infer type + human_description: + en_US: tts infer type + zh_Hans: 合成方式 + pt_BR: tts infer type + llm_description: tts infer type + options: + - value: PresetVoice + label: + en_US: preset voice + zh_Hans: 预置音色 + - value: CloneVoice + label: + en_US: clone voice + zh_Hans: 克隆音色 + - value: CloneVoice_CrossLingual + label: + en_US: clone crossLingual voice + zh_Hans: 克隆音色(跨语言) + - value: InstructVoice + label: + en_US: instruct voice + zh_Hans: 指令音色 + form: form + - name: voice + type: select + required: false + label: + en_US: preset voice + zh_Hans: 预置音色 + pt_BR: preset voice + human_description: + en_US: preset voice + zh_Hans: 预置音色 + pt_BR: preset voice + llm_description: preset voice + options: + - value: 中文男 + label: + en_US: zh-cn male + zh_Hans: 中文男 + - value: 中文女 + label: + en_US: zh-cn female + zh_Hans: 中文女 + - value: 粤语女 + label: + en_US: zh-TW female + zh_Hans: 粤语女 + form: form + - name: mock_voice_audio + type: string + required: false + label: + en_US: clone voice link + zh_Hans: 克隆音频链接 + pt_BR: clone voice link + human_description: + en_US: clone voice link + zh_Hans: 克隆音频链接 + pt_BR: clone voice link + llm_description: clone voice link + form: llm + - name: mock_voice_text + type: string + required: false + label: + en_US: text of clone voice + zh_Hans: 克隆音频对应文本 + pt_BR: text of clone voice + human_description: + en_US: text of clone voice + zh_Hans: 克隆音频对应文本 + pt_BR: text of clone voice + llm_description: text of clone voice + form: llm + - name: voice_instruct_prompt + type: string + required: false + label: + en_US: instruct prompt for voice + zh_Hans: 音色指令文本 + pt_BR: instruct prompt for voice + human_description: + en_US: instruct prompt for voice + zh_Hans: 音色指令文本 + pt_BR: instruct prompt for voice + llm_description: instruct prompt for voice + form: llm + - name: aws_region + type: string + required: false + label: + en_US: region of sagemaker endpoint + zh_Hans: SageMaker 端点所在的region + pt_BR: region of sagemaker endpoint + human_description: + en_US: region of sagemaker endpoint + zh_Hans: SageMaker 端点所在的region + pt_BR: region of sagemaker endpoint + llm_description: region of sagemaker endpoint + form: form diff --git a/api/core/tools/provider/builtin/aws/tools/transcribe_asr.py b/api/core/tools/provider/builtin/aws/tools/transcribe_asr.py new file mode 100644 index 0000000000000000000000000000000000000000..7520f6bca8b1cef4ff6276d0691b5c4da7527dcb --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/transcribe_asr.py @@ -0,0 +1,418 @@ +import json +import logging +import os +import re +import time +import uuid +from typing import Any, Union +from urllib.parse import urlparse + +import boto3 +import requests +from botocore.exceptions import ClientError +from requests.exceptions import RequestException + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + + +LanguageCodeOptions = [ + "af-ZA", + "ar-AE", + "ar-SA", + "da-DK", + "de-CH", + "de-DE", + "en-AB", + "en-AU", + "en-GB", + "en-IE", + "en-IN", + "en-US", + "en-WL", + "es-ES", + "es-US", + "fa-IR", + "fr-CA", + "fr-FR", + "he-IL", + "hi-IN", + "id-ID", + "it-IT", + "ja-JP", + "ko-KR", + "ms-MY", + "nl-NL", + "pt-BR", + "pt-PT", + "ru-RU", + "ta-IN", + "te-IN", + "tr-TR", + "zh-CN", + "zh-TW", + "th-TH", + "en-ZA", + "en-NZ", + "vi-VN", + "sv-SE", + "ab-GE", + "ast-ES", + "az-AZ", + "ba-RU", + "be-BY", + "bg-BG", + "bn-IN", + "bs-BA", + "ca-ES", + "ckb-IQ", + "ckb-IR", + "cs-CZ", + "cy-WL", + "el-GR", + "et-ET", + "eu-ES", + "fi-FI", + "gl-ES", + "gu-IN", + "ha-NG", + "hr-HR", + "hu-HU", + "hy-AM", + "is-IS", + "ka-GE", + "kab-DZ", + "kk-KZ", + "kn-IN", + "ky-KG", + "lg-IN", + "lt-LT", + "lv-LV", + "mhr-RU", + "mi-NZ", + "mk-MK", + "ml-IN", + "mn-MN", + "mr-IN", + "mt-MT", + "no-NO", + "or-IN", + "pa-IN", + "pl-PL", + "ps-AF", + "ro-RO", + "rw-RW", + "si-LK", + "sk-SK", + "sl-SI", + "so-SO", + "sr-RS", + "su-ID", + "sw-BI", + "sw-KE", + "sw-RW", + "sw-TZ", + "sw-UG", + "tl-PH", + "tt-RU", + "ug-CN", + "uk-UA", + "uz-UZ", + "wo-SN", + "zu-ZA", +] + +MediaFormat = ["mp3", "mp4", "wav", "flac", "ogg", "amr", "webm", "m4a"] + + +def is_url(text): + if not text: + return False + text = text.strip() + # Regular expression pattern for URL validation + pattern = re.compile( + r"^" # Start of the string + r"(?:http|https)://" # Protocol (http or https) + r"(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+(?:[A-Z]{2,6}\.?|[A-Z0-9-]{2,}\.?)|" # Domain + r"localhost|" # localhost + r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})" # IP address + r"(?::\d+)?" # Optional port + r"(?:/?|[/?]\S+)" # Path + r"$", # End of the string + re.IGNORECASE, + ) + return bool(pattern.match(text)) + + +def upload_file_from_url_to_s3(s3_client, url, bucket_name, s3_key=None, max_retries=3): + """ + Upload a file from a URL to an S3 bucket with retries and better error handling. + + Parameters: + - s3_client + - url (str): The URL of the file to upload + - bucket_name (str): The name of the S3 bucket + - s3_key (str): The desired key (path) in S3. If None, will use the filename from URL + - max_retries (int): Maximum number of retry attempts + + Returns: + - tuple: (bool, str) - (Success status, Message) + """ + + # Validate inputs + if not url or not bucket_name: + return False, "URL and bucket name are required" + + retry_count = 0 + while retry_count < max_retries: + try: + # Download the file from URL + response = requests.get(url, stream=True, timeout=30) + response.raise_for_status() + + # If s3_key is not provided, try to get filename from URL + if not s3_key: + parsed_url = urlparse(url) + filename = os.path.basename(parsed_url.path.split("/file-preview")[0]) + s3_key = "transcribe-files/" + filename + + # Upload the file to S3 + s3_client.upload_fileobj( + response.raw, + bucket_name, + s3_key, + ExtraArgs={ + "ContentType": response.headers.get("content-type"), + "ACL": "private", # Ensure the uploaded file is private + }, + ) + + return f"s3://{bucket_name}/{s3_key}", f"Successfully uploaded file to s3://{bucket_name}/{s3_key}" + + except RequestException as e: + retry_count += 1 + if retry_count == max_retries: + return None, f"Failed to download file from URL after {max_retries} attempts: {str(e)}" + continue + + except ClientError as e: + return None, f"AWS S3 error: {str(e)}" + + except Exception as e: + return None, f"Unexpected error: {str(e)}" + + return None, "Maximum retries exceeded" + + +class TranscribeTool(BuiltinTool): + s3_client: Any = None + transcribe_client: Any = None + + """ + Note that you must include one of LanguageCode, IdentifyLanguage, + or IdentifyMultipleLanguages in your request. + If you include more than one of these parameters, your transcription job fails. + """ + + def _transcribe_audio(self, audio_file_uri, file_type, **extra_args): + uuid_str = str(uuid.uuid4()) + job_name = f"{int(time.time())}-{uuid_str}" + try: + # Start transcription job + response = self.transcribe_client.start_transcription_job( + TranscriptionJobName=job_name, Media={"MediaFileUri": audio_file_uri}, **extra_args + ) + + # Wait for the job to complete + while True: + status = self.transcribe_client.get_transcription_job(TranscriptionJobName=job_name) + if status["TranscriptionJob"]["TranscriptionJobStatus"] in ["COMPLETED", "FAILED"]: + break + time.sleep(5) + + if status["TranscriptionJob"]["TranscriptionJobStatus"] == "COMPLETED": + return status["TranscriptionJob"]["Transcript"]["TranscriptFileUri"], None + else: + return None, f"Error: TranscriptionJobStatus:{status['TranscriptionJob']['TranscriptionJobStatus']} " + + except Exception as e: + return None, f"Error: {str(e)}" + + def _download_and_read_transcript(self, transcript_file_uri: str, max_retries: int = 3) -> tuple[str, str]: + """ + Download and read the transcript file from the given URI. + + Parameters: + - transcript_file_uri (str): The URI of the transcript file + - max_retries (int): Maximum number of retry attempts + + Returns: + - tuple: (text, error) - (Transcribed text if successful, error message if failed) + """ + retry_count = 0 + while retry_count < max_retries: + try: + # Download the transcript file + response = requests.get(transcript_file_uri, timeout=30) + response.raise_for_status() + + # Parse the JSON content + transcript_data = response.json() + + # Check if speaker labels are present and enabled + has_speaker_labels = ( + "results" in transcript_data + and "speaker_labels" in transcript_data["results"] + and "segments" in transcript_data["results"]["speaker_labels"] + ) + + if has_speaker_labels: + # Get speaker segments + segments = transcript_data["results"]["speaker_labels"]["segments"] + items = transcript_data["results"]["items"] + + # Create a mapping of start_time -> speaker_label + time_to_speaker = {} + for segment in segments: + speaker_label = segment["speaker_label"] + for item in segment["items"]: + time_to_speaker[item["start_time"]] = speaker_label + + # Build transcript with speaker labels + current_speaker = None + transcript_parts = [] + + for item in items: + # Skip non-pronunciation items (like punctuation) + if item["type"] == "punctuation": + transcript_parts.append(item["alternatives"][0]["content"]) + continue + + start_time = item["start_time"] + speaker = time_to_speaker.get(start_time) + + if speaker != current_speaker: + current_speaker = speaker + transcript_parts.append(f"\n[{speaker}]: ") + + transcript_parts.append(item["alternatives"][0]["content"]) + + return " ".join(transcript_parts).strip(), None + else: + # Extract the transcription text + # The transcript text is typically in the 'results' -> 'transcripts' array + if "results" in transcript_data and "transcripts" in transcript_data["results"]: + transcripts = transcript_data["results"]["transcripts"] + if transcripts: + # Combine all transcript segments + full_text = " ".join(t.get("transcript", "") for t in transcripts) + return full_text, None + + return None, "No transcripts found in the response" + + except requests.exceptions.RequestException as e: + retry_count += 1 + if retry_count == max_retries: + return None, f"Failed to download transcript file after {max_retries} attempts: {str(e)}" + continue + + except json.JSONDecodeError as e: + return None, f"Failed to parse transcript JSON: {str(e)}" + + except Exception as e: + return None, f"Unexpected error while processing transcript: {str(e)}" + + return None, "Maximum retries exceeded" + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + try: + if not self.transcribe_client: + aws_region = tool_parameters.get("aws_region") + if aws_region: + self.transcribe_client = boto3.client("transcribe", region_name=aws_region) + self.s3_client = boto3.client("s3", region_name=aws_region) + else: + self.transcribe_client = boto3.client("transcribe") + self.s3_client = boto3.client("s3") + + file_url = tool_parameters.get("file_url") + file_type = tool_parameters.get("file_type") + language_code = tool_parameters.get("language_code") + identify_language = tool_parameters.get("identify_language", True) + identify_multiple_languages = tool_parameters.get("identify_multiple_languages", False) + language_options_str = tool_parameters.get("language_options") + s3_bucket_name = tool_parameters.get("s3_bucket_name") + ShowSpeakerLabels = tool_parameters.get("ShowSpeakerLabels", True) + MaxSpeakerLabels = tool_parameters.get("MaxSpeakerLabels", 2) + + # Check the input params + if not s3_bucket_name: + return self.create_text_message(text="s3_bucket_name is required") + language_options = None + if language_options_str: + language_options = language_options_str.split("|") + for lang in language_options: + if lang not in LanguageCodeOptions: + return self.create_text_message( + text=f"{lang} is not supported, should be one of {LanguageCodeOptions}" + ) + if language_code and language_code not in LanguageCodeOptions: + err_msg = f"language_code:{language_code} is not supported, should be one of {LanguageCodeOptions}" + return self.create_text_message(text=err_msg) + + err_msg = f"identify_language:{identify_language}, \ + identify_multiple_languages:{identify_multiple_languages}, \ + Note that you must include one of LanguageCode, IdentifyLanguage, \ + or IdentifyMultipleLanguages in your request. \ + If you include more than one of these parameters, \ + your transcription job fails." + if not language_code: + if identify_language and identify_multiple_languages: + return self.create_text_message(text=err_msg) + else: + if identify_language or identify_multiple_languages: + return self.create_text_message(text=err_msg) + + extra_args = { + "IdentifyLanguage": identify_language, + "IdentifyMultipleLanguages": identify_multiple_languages, + } + if language_code: + extra_args["LanguageCode"] = language_code + if language_options: + extra_args["LanguageOptions"] = language_options + if ShowSpeakerLabels: + extra_args["Settings"] = {"ShowSpeakerLabels": ShowSpeakerLabels, "MaxSpeakerLabels": MaxSpeakerLabels} + + # upload to s3 bucket + s3_path_result, error = upload_file_from_url_to_s3(self.s3_client, url=file_url, bucket_name=s3_bucket_name) + if not s3_path_result: + return self.create_text_message(text=error) + + transcript_file_uri, error = self._transcribe_audio( + audio_file_uri=s3_path_result, + file_type=file_type, + **extra_args, + ) + if not transcript_file_uri: + return self.create_text_message(text=error) + + # Download and read the transcript + transcript_text, error = self._download_and_read_transcript(transcript_file_uri) + if not transcript_text: + return self.create_text_message(text=error) + + return self.create_text_message(text=transcript_text) + + except Exception as e: + return self.create_text_message(f"Exception {str(e)}") diff --git a/api/core/tools/provider/builtin/aws/tools/transcribe_asr.yaml b/api/core/tools/provider/builtin/aws/tools/transcribe_asr.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0dccd615d272dd08d2f86db7c7c5ed56faa39010 --- /dev/null +++ b/api/core/tools/provider/builtin/aws/tools/transcribe_asr.yaml @@ -0,0 +1,133 @@ +identity: + name: transcribe_asr + author: AWS + label: + en_US: TranscribeASR + zh_Hans: Transcribe语音识别转录 + pt_BR: TranscribeASR + icon: icon.svg +description: + human: + en_US: A tool for ASR (Automatic Speech Recognition) - https://github.com/aws-samples/dify-aws-tool + zh_Hans: AWS 语音识别转录服务, 请参考 https://aws.amazon.com/cn/pm/transcribe/#Learn_More_About_Amazon_Transcribe + pt_BR: A tool for ASR (Automatic Speech Recognition). + llm: A tool for ASR (Automatic Speech Recognition). +parameters: + - name: file_url + type: string + required: true + label: + en_US: video or audio file url for transcribe + zh_Hans: 语音或者视频文件url + pt_BR: video or audio file url for transcribe + human_description: + en_US: video or audio file url for transcribe + zh_Hans: 语音或者视频文件url + pt_BR: video or audio file url for transcribe + llm_description: video or audio file url for transcribe + form: llm + - name: language_code + type: string + required: false + label: + en_US: Language Code + zh_Hans: 语言编码 + pt_BR: Language Code + human_description: + en_US: The language code used to create your transcription job. refer to :https://docs.aws.amazon.com/transcribe/latest/dg/supported-languages.html + zh_Hans: 语言编码,例如zh-CN, en-US 可参考 https://docs.aws.amazon.com/transcribe/latest/dg/supported-languages.html + pt_BR: The language code used to create your transcription job. refer to :https://docs.aws.amazon.com/transcribe/latest/dg/supported-languages.html + llm_description: The language code used to create your transcription job. + form: llm + - name: identify_language + type: boolean + default: true + required: false + label: + en_US: Automactically Identify Language + zh_Hans: 自动识别语言 + pt_BR: Automactically Identify Language + human_description: + en_US: Automactically Identify Language + zh_Hans: 自动识别语言 + pt_BR: Automactically Identify Language + llm_description: Enable Automactically Identify Language + form: form + - name: identify_multiple_languages + type: boolean + required: false + label: + en_US: Automactically Identify Multiple Languages + zh_Hans: 自动识别多种语言 + pt_BR: Automactically Identify Multiple Languages + human_description: + en_US: Automactically Identify Multiple Languages + zh_Hans: 自动识别多种语言 + pt_BR: Automactically Identify Multiple Languages + llm_description: Enable Automactically Identify Multiple Languages + form: form + - name: language_options + type: string + required: false + label: + en_US: Language Options + zh_Hans: 语言种类选项 + pt_BR: Language Options + human_description: + en_US: Seperated by |, e.g:zh-CN|en-US, You can specify two or more language codes that represent the languages you think may be present in your media + zh_Hans: 您可以指定两个或更多的语言代码来表示您认为可能出现在媒体中的语言。用|分隔,如 zh-CN|en-US + pt_BR: Seperated by |, e.g:zh-CN|en-US, You can specify two or more language codes that represent the languages you think may be present in your media + llm_description: Seperated by |, e.g:zh-CN|en-US, You can specify two or more language codes that represent the languages you think may be present in your media + form: llm + - name: s3_bucket_name + type: string + required: true + label: + en_US: s3 bucket name + zh_Hans: s3 存储桶名称 + pt_BR: s3 bucket name + human_description: + en_US: s3 bucket name to store transcribe files (don't add prefix s3://) + zh_Hans: s3 存储桶名称,用于存储转录文件 (不需要前缀 s3://) + pt_BR: s3 bucket name to store transcribe files (don't add prefix s3://) + llm_description: s3 bucket name to store transcribe files + form: form + - name: ShowSpeakerLabels + type: boolean + required: true + default: true + label: + en_US: ShowSpeakerLabels + zh_Hans: 显示说话人标签 + pt_BR: ShowSpeakerLabels + human_description: + en_US: Enables speaker partitioning (diarization) in your transcription output + zh_Hans: 在转录输出中启用说话人分区(说话人分离) + pt_BR: Enables speaker partitioning (diarization) in your transcription output + llm_description: Enables speaker partitioning (diarization) in your transcription output + form: form + - name: MaxSpeakerLabels + type: number + required: true + default: 2 + label: + en_US: MaxSpeakerLabels + zh_Hans: 说话人标签数量 + pt_BR: MaxSpeakerLabels + human_description: + en_US: Specify the maximum number of speakers you want to partition in your media + zh_Hans: 指定您希望在媒体中划分的最多演讲者数量。 + pt_BR: Specify the maximum number of speakers you want to partition in your media + llm_description: Specify the maximum number of speakers you want to partition in your media + form: form + - name: aws_region + type: string + required: false + label: + en_US: AWS Region + zh_Hans: AWS 区域 + human_description: + en_US: Please enter the AWS region for the transcribe service, for example 'us-east-1'. + zh_Hans: 请输入Transcribe的 AWS 区域,例如 'us-east-1'。 + llm_description: Please enter the AWS region for the transcribe service, for example 'us-east-1'. + form: form diff --git a/api/core/tools/provider/builtin/azuredalle/__init__.py b/api/core/tools/provider/builtin/azuredalle/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/tools/provider/builtin/azuredalle/_assets/icon.png b/api/core/tools/provider/builtin/azuredalle/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..7083a3f638e9a18e2d9c09616bd1b9b5e36f53cb Binary files /dev/null and b/api/core/tools/provider/builtin/azuredalle/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/azuredalle/azuredalle.py b/api/core/tools/provider/builtin/azuredalle/azuredalle.py new file mode 100644 index 0000000000000000000000000000000000000000..1fab0d03a28ff3096b0618ac04f1ed7fd4b608e8 --- /dev/null +++ b/api/core/tools/provider/builtin/azuredalle/azuredalle.py @@ -0,0 +1,20 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.azuredalle.tools.dalle3 import DallE3Tool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class AzureDALLEProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + DallE3Tool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={"prompt": "cute girl, blue eyes, white hair, anime style", "size": "square", "n": 1}, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/azuredalle/azuredalle.yaml b/api/core/tools/provider/builtin/azuredalle/azuredalle.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4353e0c4862f619abe496992b21d39b400cf156f --- /dev/null +++ b/api/core/tools/provider/builtin/azuredalle/azuredalle.yaml @@ -0,0 +1,76 @@ +identity: + author: Leslie + name: azuredalle + label: + en_US: Azure DALL-E + zh_Hans: Azure DALL-E 绘画 + pt_BR: Azure DALL-E + description: + en_US: Azure DALL-E art + zh_Hans: Azure DALL-E 绘画 + pt_BR: Azure DALL-E art + icon: icon.png + tags: + - image + - productivity +credentials_for_provider: + azure_openai_api_key: + type: secret-input + required: true + label: + en_US: API key + zh_Hans: 密钥 + pt_BR: API key + help: + en_US: Please input your Azure OpenAI API key + zh_Hans: 请输入你的 Azure OpenAI API key + pt_BR: Introduza a sua chave de API OpenAI do Azure + placeholder: + en_US: Please input your Azure OpenAI API key + zh_Hans: 请输入你的 Azure OpenAI API key + pt_BR: Introduza a sua chave de API OpenAI do Azure + azure_openai_api_model_name: + type: text-input + required: true + label: + en_US: Deployment Name + zh_Hans: 部署名称 + pt_BR: Nome da Implantação + help: + en_US: Please input the name of your Azure Openai DALL-E API deployment + zh_Hans: 请输入你的 Azure Openai DALL-E API 部署名称 + pt_BR: Insira o nome da implantação da API DALL-E do Azure Openai + placeholder: + en_US: Please input the name of your Azure Openai DALL-E API deployment + zh_Hans: 请输入你的 Azure Openai DALL-E API 部署名称 + pt_BR: Insira o nome da implantação da API DALL-E do Azure Openai + azure_openai_base_url: + type: text-input + required: true + label: + en_US: API Endpoint URL + zh_Hans: API 域名 + pt_BR: API Endpoint URL + help: + en_US: Please input your Azure OpenAI Endpoint URL, e.g. https://xxx.openai.azure.com/ + zh_Hans: 请输入你的 Azure OpenAI API域名,例如:https://xxx.openai.azure.com/ + pt_BR: Introduza a URL do Azure OpenAI Endpoint, e.g. https://xxx.openai.azure.com/ + placeholder: + en_US: Please input your Azure OpenAI Endpoint URL, e.g. https://xxx.openai.azure.com/ + zh_Hans: 请输入你的 Azure OpenAI API域名,例如:https://xxx.openai.azure.com/ + pt_BR: Introduza a URL do Azure OpenAI Endpoint, e.g. https://xxx.openai.azure.com/ + azure_openai_api_version: + type: text-input + required: true + label: + en_US: API Version + zh_Hans: API 版本 + pt_BR: API Version + help: + en_US: Please input your Azure OpenAI API Version,e.g. 2023-12-01-preview + zh_Hans: 请输入你的 Azure OpenAI API 版本,例如:2023-12-01-preview + pt_BR: Introduza a versão da API OpenAI do Azure,e.g. 2023-12-01-preview + placeholder: + en_US: Please input your Azure OpenAI API Version,e.g. 2023-12-01-preview + zh_Hans: 请输入你的 Azure OpenAI API 版本,例如:2023-12-01-preview + pt_BR: Introduza a versão da API OpenAI do Azure,e.g. 2023-12-01-preview diff --git a/api/core/tools/provider/builtin/azuredalle/tools/dalle3.py b/api/core/tools/provider/builtin/azuredalle/tools/dalle3.py new file mode 100644 index 0000000000000000000000000000000000000000..cfa3cfb092803a5ad3807eb348c793237d7643ed --- /dev/null +++ b/api/core/tools/provider/builtin/azuredalle/tools/dalle3.py @@ -0,0 +1,83 @@ +import random +from base64 import b64decode +from typing import Any, Union + +from openai import AzureOpenAI + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class DallE3Tool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + client = AzureOpenAI( + api_version=self.runtime.credentials["azure_openai_api_version"], + azure_endpoint=self.runtime.credentials["azure_openai_base_url"], + api_key=self.runtime.credentials["azure_openai_api_key"], + ) + + SIZE_MAPPING = { + "square": "1024x1024", + "vertical": "1024x1792", + "horizontal": "1792x1024", + } + + # prompt + prompt = tool_parameters.get("prompt", "") + if not prompt: + return self.create_text_message("Please input prompt") + # get size + size = SIZE_MAPPING[tool_parameters.get("size", "square")] + # get n + n = tool_parameters.get("n", 1) + # get quality + quality = tool_parameters.get("quality", "standard") + if quality not in {"standard", "hd"}: + return self.create_text_message("Invalid quality") + # get style + style = tool_parameters.get("style", "vivid") + if style not in {"natural", "vivid"}: + return self.create_text_message("Invalid style") + # set extra body + seed_id = tool_parameters.get("seed_id", self._generate_random_id(8)) + extra_body = {"seed": seed_id} + + # call openapi dalle3 + model = self.runtime.credentials["azure_openai_api_model_name"] + response = client.images.generate( + prompt=prompt, + model=model, + size=size, + n=n, + extra_body=extra_body, + style=style, + quality=quality, + response_format="b64_json", + ) + + result = [] + + for image in response.data: + result.append( + self.create_blob_message( + blob=b64decode(image.b64_json), + meta={"mime_type": "image/png"}, + save_as=self.VariableKey.IMAGE.value, + ) + ) + result.append(self.create_text_message(f"\nGenerate image source to Seed ID: {seed_id}")) + + return result + + @staticmethod + def _generate_random_id(length=8): + characters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789" + random_id = "".join(random.choices(characters, k=length)) + return random_id diff --git a/api/core/tools/provider/builtin/azuredalle/tools/dalle3.yaml b/api/core/tools/provider/builtin/azuredalle/tools/dalle3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e256748e8f718880bf43a3d457d6cc7caa7af3ac --- /dev/null +++ b/api/core/tools/provider/builtin/azuredalle/tools/dalle3.yaml @@ -0,0 +1,136 @@ +identity: + name: azure_dalle3 + author: Leslie + label: + en_US: Azure DALL-E 3 + zh_Hans: Azure DALL-E 3 绘画 + pt_BR: Azure DALL-E 3 + description: + en_US: DALL-E 3 is a powerful drawing tool that can draw the image you want based on your prompt, compared to DallE 2, DallE 3 has stronger drawing ability, but it will consume more resources + zh_Hans: DALL-E 3 是一个强大的绘画工具,它可以根据您的提示词绘制出您想要的图像,相比于DallE 2, DallE 3拥有更强的绘画能力,但会消耗更多的资源 + pt_BR: DALL-E 3 é uma poderosa ferramenta de desenho que pode desenhar a imagem que você deseja com base em seu prompt, em comparação com DallE 2, DallE 3 tem uma capacidade de desenho mais forte, mas consumirá mais recursos +description: + human: + en_US: DALL-E is a text to image tool + zh_Hans: DALL-E 是一个文本到图像的工具 + pt_BR: DALL-E é uma ferramenta de texto para imagem + llm: DALL-E is a tool used to generate images from text +parameters: + - name: prompt + type: string + required: true + label: + en_US: Prompt + zh_Hans: 提示词 + pt_BR: Prompt + human_description: + en_US: Image prompt, you can check the official documentation of DallE 3 + zh_Hans: 图像提示词,您可以查看 DallE 3 的官方文档 + pt_BR: Imagem prompt, você pode verificar a documentação oficial do DallE 3 + llm_description: Image prompt of DallE 3, you should describe the image you want to generate as a list of words as possible as detailed + form: llm + - name: seed_id + type: string + required: false + label: + en_US: Seed ID + zh_Hans: 种子ID + pt_BR: ID da semente + human_description: + en_US: Image generation seed ID to ensure consistency of series generated images + zh_Hans: 图像生成种子ID,确保系列生成图像的一致性 + pt_BR: ID de semente de geração de imagem para garantir a consistência das imagens geradas em série + llm_description: If the user requests image consistency, extract the seed ID from the user's question or context.The seed id consists of an 8-bit string containing uppercase and lowercase letters and numbers + form: llm + - name: size + type: select + required: true + human_description: + en_US: selecting the image size + zh_Hans: 选择图像大小 + pt_BR: seleccionar o tamanho da imagem + label: + en_US: Image size + zh_Hans: 图像大小 + pt_BR: Tamanho da imagem + form: form + options: + - value: square + label: + en_US: Squre(1024x1024) + zh_Hans: 方(1024x1024) + pt_BR: Squire(1024x1024) + - value: vertical + label: + en_US: Vertical(1024x1792) + zh_Hans: 竖屏(1024x1792) + pt_BR: Vertical(1024x1792) + - value: horizontal + label: + en_US: Horizontal(1792x1024) + zh_Hans: 横屏(1792x1024) + pt_BR: Horizontal(1792x1024) + default: square + - name: n + type: number + required: true + human_description: + en_US: selecting the number of images + zh_Hans: 选择图像数量 + pt_BR: seleccionar o número de imagens + label: + en_US: Number of images + zh_Hans: 图像数量 + pt_BR: Número de imagens + form: form + min: 1 + max: 1 + default: 1 + - name: quality + type: select + required: true + human_description: + en_US: selecting the image quality + zh_Hans: 选择图像质量 + pt_BR: seleccionar a qualidade da imagem + label: + en_US: Image quality + zh_Hans: 图像质量 + pt_BR: Qualidade da imagem + form: form + options: + - value: standard + label: + en_US: Standard + zh_Hans: 标准 + pt_BR: Normal + - value: hd + label: + en_US: HD + zh_Hans: 高清 + pt_BR: HD + default: standard + - name: style + type: select + required: true + human_description: + en_US: selecting the image style + zh_Hans: 选择图像风格 + pt_BR: seleccionar o estilo da imagem + label: + en_US: Image style + zh_Hans: 图像风格 + pt_BR: Estilo da imagem + form: form + options: + - value: vivid + label: + en_US: Vivid + zh_Hans: 生动 + pt_BR: Vívido + - value: natural + label: + en_US: Natural + zh_Hans: 自然 + pt_BR: Natural + default: vivid diff --git a/api/core/tools/provider/builtin/baidu_translate/_assets/icon.png b/api/core/tools/provider/builtin/baidu_translate/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..8eb8f21513ba7d45de8204bfe64aa3cc1fd7fc26 Binary files /dev/null and b/api/core/tools/provider/builtin/baidu_translate/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/baidu_translate/_baidu_translate_tool_base.py b/api/core/tools/provider/builtin/baidu_translate/_baidu_translate_tool_base.py new file mode 100644 index 0000000000000000000000000000000000000000..ce907c3c616e07d5356359c482c79c24396b1bf5 --- /dev/null +++ b/api/core/tools/provider/builtin/baidu_translate/_baidu_translate_tool_base.py @@ -0,0 +1,11 @@ +from hashlib import md5 + + +class BaiduTranslateToolBase: + def _get_sign(self, appid, secret, salt, query): + """ + get baidu translate sign + """ + # concatenate the string in the order of appid+q+salt+secret + str = appid + query + salt + secret + return md5(str.encode("utf-8")).hexdigest() diff --git a/api/core/tools/provider/builtin/baidu_translate/baidu_translate.py b/api/core/tools/provider/builtin/baidu_translate/baidu_translate.py new file mode 100644 index 0000000000000000000000000000000000000000..cccd2f8c8fc4786eb9b3a43d3258c4f9437df521 --- /dev/null +++ b/api/core/tools/provider/builtin/baidu_translate/baidu_translate.py @@ -0,0 +1,17 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.baidu_translate.tools.translate import BaiduTranslateTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class BaiduTranslateProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + BaiduTranslateTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke(user_id="", tool_parameters={"q": "这是一段测试文本", "from": "auto", "to": "en"}) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/baidu_translate/baidu_translate.yaml b/api/core/tools/provider/builtin/baidu_translate/baidu_translate.yaml new file mode 100644 index 0000000000000000000000000000000000000000..06dadeeefc9cded48bdea389c910a43bc12b9518 --- /dev/null +++ b/api/core/tools/provider/builtin/baidu_translate/baidu_translate.yaml @@ -0,0 +1,39 @@ +identity: + author: Xiao Ley + name: baidu_translate + label: + en_US: Baidu Translate + zh_Hans: 百度翻译 + description: + en_US: Translate text using Baidu + zh_Hans: 使用百度进行翻译 + icon: icon.png + tags: + - utilities +credentials_for_provider: + appid: + type: secret-input + required: true + label: + en_US: Baidu translate appid + zh_Hans: Baidu translate appid + placeholder: + en_US: Please input your Baidu translate appid + zh_Hans: 请输入你的百度翻译 appid + help: + en_US: Get your Baidu translate appid from Baidu translate + zh_Hans: 从百度翻译开放平台获取你的 appid + url: https://api.fanyi.baidu.com + secret: + type: secret-input + required: true + label: + en_US: Baidu translate secret + zh_Hans: Baidu translate secret + placeholder: + en_US: Please input your Baidu translate secret + zh_Hans: 请输入你的百度翻译 secret + help: + en_US: Get your Baidu translate secret from Baidu translate + zh_Hans: 从百度翻译开放平台获取你的 secret + url: https://api.fanyi.baidu.com diff --git a/api/core/tools/provider/builtin/baidu_translate/tools/fieldtranslate.py b/api/core/tools/provider/builtin/baidu_translate/tools/fieldtranslate.py new file mode 100644 index 0000000000000000000000000000000000000000..ff5cf32ddc7525544bdeca04203717a544bc9cae --- /dev/null +++ b/api/core/tools/provider/builtin/baidu_translate/tools/fieldtranslate.py @@ -0,0 +1,78 @@ +import random +from hashlib import md5 +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.baidu_translate._baidu_translate_tool_base import BaiduTranslateToolBase +from core.tools.tool.builtin_tool import BuiltinTool + + +class BaiduFieldTranslateTool(BuiltinTool, BaiduTranslateToolBase): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + BAIDU_FIELD_TRANSLATE_URL = "https://fanyi-api.baidu.com/api/trans/vip/fieldtranslate" + + appid = self.runtime.credentials.get("appid", "") + if not appid: + raise ValueError("invalid baidu translate appid") + + secret = self.runtime.credentials.get("secret", "") + if not secret: + raise ValueError("invalid baidu translate secret") + + q = tool_parameters.get("q", "") + if not q: + raise ValueError("Please input text to translate") + + from_ = tool_parameters.get("from", "") + if not from_: + raise ValueError("Please select source language") + + to = tool_parameters.get("to", "") + if not to: + raise ValueError("Please select destination language") + + domain = tool_parameters.get("domain", "") + if not domain: + raise ValueError("Please select domain") + + salt = str(random.randint(32768, 16777215)) + sign = self._get_sign(appid, secret, salt, q, domain) + + headers = {"Content-Type": "application/x-www-form-urlencoded"} + params = { + "q": q, + "from": from_, + "to": to, + "appid": appid, + "salt": salt, + "domain": domain, + "sign": sign, + "needIntervene": 1, + } + try: + response = requests.post(BAIDU_FIELD_TRANSLATE_URL, headers=headers, data=params) + result = response.json() + + if "trans_result" in result: + result_text = result["trans_result"][0]["dst"] + else: + result_text = f"{result['error_code']}: {result['error_msg']}" + + return self.create_text_message(str(result_text)) + except requests.RequestException as e: + raise ValueError(f"Translation service error: {e}") + except Exception: + raise ValueError("Translation service error, please check the network") + + def _get_sign(self, appid, secret, salt, query, domain): + str = appid + query + salt + domain + secret + return md5(str.encode("utf-8")).hexdigest() diff --git a/api/core/tools/provider/builtin/baidu_translate/tools/fieldtranslate.yaml b/api/core/tools/provider/builtin/baidu_translate/tools/fieldtranslate.yaml new file mode 100644 index 0000000000000000000000000000000000000000..de51fddbaea42259a570df0ec810e2857c1e0ee2 --- /dev/null +++ b/api/core/tools/provider/builtin/baidu_translate/tools/fieldtranslate.yaml @@ -0,0 +1,123 @@ +identity: + name: field_translate + author: Xiao Ley + label: + en_US: Field translate + zh_Hans: 百度领域翻译 +description: + human: + en_US: A tool for Baidu Field translate (Currently, the fields of "novel" and "wiki" only support Chinese to English translation. If the language direction is set to English to Chinese, the default output will be a universal translation result). + zh_Hans: 百度领域翻译,提供多种领域的文本翻译(目前“网络文学领域”和“人文社科领域”仅支持中到英,如设置语言方向为英到中,则默认输出通用翻译结果) + llm: A tool for Baidu Field translate +parameters: + - name: q + type: string + required: true + label: + en_US: Text content + zh_Hans: 文本内容 + human_description: + en_US: Text content to be translated + zh_Hans: 需要翻译的文本内容 + llm_description: Text content to be translated + form: llm + - name: from + type: select + required: true + label: + en_US: source language + zh_Hans: 源语言 + human_description: + en_US: The source language of the input text + zh_Hans: 输入的文本的源语言 + default: auto + form: form + options: + - value: auto + label: + en_US: auto + zh_Hans: 自动检测 + - value: zh + label: + en_US: Chinese + zh_Hans: 中文 + - value: en + label: + en_US: English + zh_Hans: 英语 + - name: to + type: select + required: true + label: + en_US: destination language + zh_Hans: 目标语言 + human_description: + en_US: The destination language of the input text + zh_Hans: 输入文本的目标语言 + default: en + form: form + options: + - value: zh + label: + en_US: Chinese + zh_Hans: 中文 + - value: en + label: + en_US: English + zh_Hans: 英语 + - name: domain + type: select + required: true + label: + en_US: domain + zh_Hans: 领域 + human_description: + en_US: The domain of the input text + zh_Hans: 输入文本的领域 + default: novel + form: form + options: + - value: it + label: + en_US: it + zh_Hans: 信息技术领域 + - value: finance + label: + en_US: finance + zh_Hans: 金融财经领域 + - value: machinery + label: + en_US: machinery + zh_Hans: 机械制造领域 + - value: senimed + label: + en_US: senimed + zh_Hans: 生物医药领域 + - value: novel + label: + en_US: novel (only support Chinese to English translation) + zh_Hans: 网络文学领域(仅支持中到英) + - value: academic + label: + en_US: academic + zh_Hans: 学术论文领域 + - value: aerospace + label: + en_US: aerospace + zh_Hans: 航空航天领域 + - value: wiki + label: + en_US: wiki (only support Chinese to English translation) + zh_Hans: 人文社科领域(仅支持中到英) + - value: news + label: + en_US: news + zh_Hans: 新闻咨询领域 + - value: law + label: + en_US: law + zh_Hans: 法律法规领域 + - value: contract + label: + en_US: contract + zh_Hans: 合同领域 diff --git a/api/core/tools/provider/builtin/baidu_translate/tools/language.py b/api/core/tools/provider/builtin/baidu_translate/tools/language.py new file mode 100644 index 0000000000000000000000000000000000000000..b7fd692b7d19046e236b1f3f9ab1de769a9b8761 --- /dev/null +++ b/api/core/tools/provider/builtin/baidu_translate/tools/language.py @@ -0,0 +1,95 @@ +import random +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.baidu_translate._baidu_translate_tool_base import BaiduTranslateToolBase +from core.tools.tool.builtin_tool import BuiltinTool + + +class BaiduLanguageTool(BuiltinTool, BaiduTranslateToolBase): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + BAIDU_LANGUAGE_URL = "https://fanyi-api.baidu.com/api/trans/vip/language" + + appid = self.runtime.credentials.get("appid", "") + if not appid: + raise ValueError("invalid baidu translate appid") + + secret = self.runtime.credentials.get("secret", "") + if not secret: + raise ValueError("invalid baidu translate secret") + + q = tool_parameters.get("q", "") + if not q: + raise ValueError("Please input text to translate") + + description_language = tool_parameters.get("description_language", "English") + + salt = str(random.randint(32768, 16777215)) + sign = self._get_sign(appid, secret, salt, q) + + headers = {"Content-Type": "application/x-www-form-urlencoded"} + params = { + "q": q, + "appid": appid, + "salt": salt, + "sign": sign, + } + + try: + response = requests.post(BAIDU_LANGUAGE_URL, params=params, headers=headers) + result = response.json() + if "error_code" not in result: + raise ValueError("Translation service error, please check the network") + + result_text = "" + if result["error_code"] != 0: + result_text = f"{result['error_code']}: {result['error_msg']}" + else: + result_text = result["data"]["src"] + result_text = self.mapping_result(description_language, result_text) + + return self.create_text_message(result_text) + except requests.RequestException as e: + raise ValueError(f"Translation service error: {e}") + except Exception: + raise ValueError("Translation service error, please check the network") + + def mapping_result(self, description_language: str, result: str) -> str: + """ + mapping result + """ + mapping = { + "English": { + "zh": "Chinese", + "en": "English", + "jp": "Japanese", + "kor": "Korean", + "th": "Thai", + "vie": "Vietnamese", + "ru": "Russian", + }, + "Chinese": { + "zh": "中文", + "en": "英文", + "jp": "日文", + "kor": "韩文", + "th": "泰语", + "vie": "越南语", + "ru": "俄语", + }, + } + + language_mapping = mapping.get(description_language) + if not language_mapping: + return result + + return language_mapping.get(result, result) diff --git a/api/core/tools/provider/builtin/baidu_translate/tools/language.yaml b/api/core/tools/provider/builtin/baidu_translate/tools/language.yaml new file mode 100644 index 0000000000000000000000000000000000000000..60cca2e288a622ca81b93169d92a48ee621127ff --- /dev/null +++ b/api/core/tools/provider/builtin/baidu_translate/tools/language.yaml @@ -0,0 +1,43 @@ +identity: + name: language + author: Xiao Ley + label: + en_US: Baidu Language + zh_Hans: 百度语种识别 +description: + human: + en_US: A tool for Baidu Language, support Chinese, English, Japanese, Korean, Thai, Vietnamese and Russian + zh_Hans: 使用百度进行语种识别,支持的语种:中文、英语、日语、韩语、泰语、越南语和俄语 + llm: A tool for Baidu Language +parameters: + - name: q + type: string + required: true + label: + en_US: Text content + zh_Hans: 文本内容 + human_description: + en_US: Text content to be recognized + zh_Hans: 需要识别语言的文本内容 + llm_description: Text content to be recognized + form: llm + - name: description_language + type: select + required: true + label: + en_US: Description language + zh_Hans: 描述语言 + human_description: + en_US: Describe the language used to identify the results + zh_Hans: 描述识别结果所用的语言 + default: Chinese + form: form + options: + - value: Chinese + label: + en_US: Chinese + zh_Hans: 中文 + - value: English + label: + en_US: English + zh_Hans: 英语 diff --git a/api/core/tools/provider/builtin/baidu_translate/tools/translate.py b/api/core/tools/provider/builtin/baidu_translate/tools/translate.py new file mode 100644 index 0000000000000000000000000000000000000000..0d25466a7060fcd1c4b4237fab541738103df265 --- /dev/null +++ b/api/core/tools/provider/builtin/baidu_translate/tools/translate.py @@ -0,0 +1,67 @@ +import random +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.baidu_translate._baidu_translate_tool_base import BaiduTranslateToolBase +from core.tools.tool.builtin_tool import BuiltinTool + + +class BaiduTranslateTool(BuiltinTool, BaiduTranslateToolBase): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + BAIDU_TRANSLATE_URL = "https://fanyi-api.baidu.com/api/trans/vip/translate" + + appid = self.runtime.credentials.get("appid", "") + if not appid: + raise ValueError("invalid baidu translate appid") + + secret = self.runtime.credentials.get("secret", "") + if not secret: + raise ValueError("invalid baidu translate secret") + + q = tool_parameters.get("q", "") + if not q: + raise ValueError("Please input text to translate") + + from_ = tool_parameters.get("from", "") + if not from_: + raise ValueError("Please select source language") + + to = tool_parameters.get("to", "") + if not to: + raise ValueError("Please select destination language") + + salt = str(random.randint(32768, 16777215)) + sign = self._get_sign(appid, secret, salt, q) + + headers = {"Content-Type": "application/x-www-form-urlencoded"} + params = { + "q": q, + "from": from_, + "to": to, + "appid": appid, + "salt": salt, + "sign": sign, + } + try: + response = requests.post(BAIDU_TRANSLATE_URL, params=params, headers=headers) + result = response.json() + + if "trans_result" in result: + result_text = result["trans_result"][0]["dst"] + else: + result_text = f"{result['error_code']}: {result['error_msg']}" + + return self.create_text_message(str(result_text)) + except requests.RequestException as e: + raise ValueError(f"Translation service error: {e}") + except Exception: + raise ValueError("Translation service error, please check the network") diff --git a/api/core/tools/provider/builtin/baidu_translate/tools/translate.yaml b/api/core/tools/provider/builtin/baidu_translate/tools/translate.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c8ff32cb6bb1f1615adf3dbd6a867ef8aace1316 --- /dev/null +++ b/api/core/tools/provider/builtin/baidu_translate/tools/translate.yaml @@ -0,0 +1,275 @@ +identity: + name: translate + author: Xiao Ley + label: + en_US: Translate + zh_Hans: 百度翻译 +description: + human: + en_US: A tool for Baidu Translate + zh_Hans: 百度翻译 + llm: A tool for Baidu Translate +parameters: + - name: q + type: string + required: true + label: + en_US: Text content + zh_Hans: 文本内容 + human_description: + en_US: Text content to be translated + zh_Hans: 需要翻译的文本内容 + llm_description: Text content to be translated + form: llm + - name: from + type: select + required: true + label: + en_US: source language + zh_Hans: 源语言 + human_description: + en_US: The source language of the input text + zh_Hans: 输入的文本的源语言 + default: auto + form: form + options: + - value: auto + label: + en_US: auto + zh_Hans: 自动检测 + - value: zh + label: + en_US: Chinese + zh_Hans: 中文 + - value: en + label: + en_US: English + zh_Hans: 英语 + - value: cht + label: + en_US: Traditional Chinese + zh_Hans: 繁体中文 + - value: yue + label: + en_US: Yue + zh_Hans: 粤语 + - value: wyw + label: + en_US: Wyw + zh_Hans: 文言文 + - value: jp + label: + en_US: Japanese + zh_Hans: 日语 + - value: kor + label: + en_US: Korean + zh_Hans: 韩语 + - value: fra + label: + en_US: French + zh_Hans: 法语 + - value: spa + label: + en_US: Spanish + zh_Hans: 西班牙语 + - value: th + label: + en_US: Thai + zh_Hans: 泰语 + - value: ara + label: + en_US: Arabic + zh_Hans: 阿拉伯语 + - value: ru + label: + en_US: Russian + zh_Hans: 俄语 + - value: pt + label: + en_US: Portuguese + zh_Hans: 葡萄牙语 + - value: de + label: + en_US: German + zh_Hans: 德语 + - value: it + label: + en_US: Italian + zh_Hans: 意大利语 + - value: el + label: + en_US: Greek + zh_Hans: 希腊语 + - value: nl + label: + en_US: Dutch + zh_Hans: 荷兰语 + - value: pl + label: + en_US: Polish + zh_Hans: 波兰语 + - value: bul + label: + en_US: Bulgarian + zh_Hans: 保加利亚语 + - value: est + label: + en_US: Estonian + zh_Hans: 爱沙尼亚语 + - value: dan + label: + en_US: Danish + zh_Hans: 丹麦语 + - value: fin + label: + en_US: Finnish + zh_Hans: 芬兰语 + - value: cs + label: + en_US: Czech + zh_Hans: 捷克语 + - value: rom + label: + en_US: Romanian + zh_Hans: 罗马尼亚语 + - value: slo + label: + en_US: Slovak + zh_Hans: 斯洛文尼亚语 + - value: swe + label: + en_US: Swedish + zh_Hans: 瑞典语 + - value: hu + label: + en_US: Hungarian + zh_Hans: 匈牙利语 + - value: vie + label: + en_US: Vietnamese + zh_Hans: 越南语 + - name: to + type: select + required: true + label: + en_US: destination language + zh_Hans: 目标语言 + human_description: + en_US: The destination language of the input text + zh_Hans: 输入文本的目标语言 + default: en + form: form + options: + - value: zh + label: + en_US: Chinese + zh_Hans: 中文 + - value: en + label: + en_US: English + zh_Hans: 英语 + - value: cht + label: + en_US: Traditional Chinese + zh_Hans: 繁体中文 + - value: yue + label: + en_US: Yue + zh_Hans: 粤语 + - value: wyw + label: + en_US: Wyw + zh_Hans: 文言文 + - value: jp + label: + en_US: Japanese + zh_Hans: 日语 + - value: kor + label: + en_US: Korean + zh_Hans: 韩语 + - value: fra + label: + en_US: French + zh_Hans: 法语 + - value: spa + label: + en_US: Spanish + zh_Hans: 西班牙语 + - value: th + label: + en_US: Thai + zh_Hans: 泰语 + - value: ara + label: + en_US: Arabic + zh_Hans: 阿拉伯语 + - value: ru + label: + en_US: Russian + zh_Hans: 俄语 + - value: pt + label: + en_US: Portuguese + zh_Hans: 葡萄牙语 + - value: de + label: + en_US: German + zh_Hans: 德语 + - value: it + label: + en_US: Italian + zh_Hans: 意大利语 + - value: el + label: + en_US: Greek + zh_Hans: 希腊语 + - value: nl + label: + en_US: Dutch + zh_Hans: 荷兰语 + - value: pl + label: + en_US: Polish + zh_Hans: 波兰语 + - value: bul + label: + en_US: Bulgarian + zh_Hans: 保加利亚语 + - value: est + label: + en_US: Estonian + zh_Hans: 爱沙尼亚语 + - value: dan + label: + en_US: Danish + zh_Hans: 丹麦语 + - value: fin + label: + en_US: Finnish + zh_Hans: 芬兰语 + - value: cs + label: + en_US: Czech + zh_Hans: 捷克语 + - value: rom + label: + en_US: Romanian + zh_Hans: 罗马尼亚语 + - value: slo + label: + en_US: Slovak + zh_Hans: 斯洛文尼亚语 + - value: swe + label: + en_US: Swedish + zh_Hans: 瑞典语 + - value: hu + label: + en_US: Hungarian + zh_Hans: 匈牙利语 + - value: vie + label: + en_US: Vietnamese + zh_Hans: 越南语 diff --git a/api/core/tools/provider/builtin/bing/_assets/icon.svg b/api/core/tools/provider/builtin/bing/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..a94de7971d35b7ba9b722941b8e294c9d9c4e304 --- /dev/null +++ b/api/core/tools/provider/builtin/bing/_assets/icon.svg @@ -0,0 +1,40 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/bing/bing.py b/api/core/tools/provider/builtin/bing/bing.py new file mode 100644 index 0000000000000000000000000000000000000000..c71128be4a784f21a6122b2a65be1e1373023323 --- /dev/null +++ b/api/core/tools/provider/builtin/bing/bing.py @@ -0,0 +1,23 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.bing.tools.bing_web_search import BingSearchTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class BingProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + BingSearchTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).validate_credentials( + credentials=credentials, + tool_parameters={ + "query": "test", + "result_type": "link", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/bing/bing.yaml b/api/core/tools/provider/builtin/bing/bing.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1ab17d5294b37c5d4f3f5eec45228f421b601771 --- /dev/null +++ b/api/core/tools/provider/builtin/bing/bing.yaml @@ -0,0 +1,107 @@ +identity: + author: Dify + name: bing + label: + en_US: Bing + zh_Hans: Bing + pt_BR: Bing + description: + en_US: Bing Search + zh_Hans: Bing 搜索 + pt_BR: Bing Search + icon: icon.svg + tags: + - search +credentials_for_provider: + subscription_key: + type: secret-input + required: true + label: + en_US: Bing subscription key + zh_Hans: Bing subscription key + pt_BR: Bing subscription key + placeholder: + en_US: Please input your Bing subscription key + zh_Hans: 请输入你的 Bing subscription key + pt_BR: Please input your Bing subscription key + help: + en_US: Get your Bing subscription key from Bing + zh_Hans: 从 Bing 获取您的 Bing subscription key + pt_BR: Get your Bing subscription key from Bing + url: https://www.microsoft.com/cognitive-services/en-us/bing-web-search-api + server_url: + type: text-input + required: false + label: + en_US: Bing endpoint + zh_Hans: Bing endpoint + pt_BR: Bing endpoint + placeholder: + en_US: Please input your Bing endpoint + zh_Hans: 请输入你的 Bing 端点 + pt_BR: Please input your Bing endpoint + help: + en_US: An endpoint is like "https://api.bing.microsoft.com/v7.0/search" + zh_Hans: 例如 "https://api.bing.microsoft.com/v7.0/search" + pt_BR: An endpoint is like "https://api.bing.microsoft.com/v7.0/search" + default: https://api.bing.microsoft.com/v7.0/search + allow_entities: + type: boolean + required: false + label: + en_US: Allow Entities Search + zh_Hans: 支持实体搜索 + pt_BR: Allow Entities Search + help: + en_US: Does your subscription plan allow entity search + zh_Hans: 您的订阅计划是否支持实体搜索 + pt_BR: Does your subscription plan allow entity search + default: true + allow_web_pages: + type: boolean + required: false + label: + en_US: Allow Web Pages Search + zh_Hans: 支持网页搜索 + pt_BR: Allow Web Pages Search + help: + en_US: Does your subscription plan allow web pages search + zh_Hans: 您的订阅计划是否支持网页搜索 + pt_BR: Does your subscription plan allow web pages search + default: true + allow_computation: + type: boolean + required: false + label: + en_US: Allow Computation Search + zh_Hans: 支持计算搜索 + pt_BR: Allow Computation Search + help: + en_US: Does your subscription plan allow computation search + zh_Hans: 您的订阅计划是否支持计算搜索 + pt_BR: Does your subscription plan allow computation search + default: false + allow_news: + type: boolean + required: false + label: + en_US: Allow News Search + zh_Hans: 支持新闻搜索 + pt_BR: Allow News Search + help: + en_US: Does your subscription plan allow news search + zh_Hans: 您的订阅计划是否支持新闻搜索 + pt_BR: Does your subscription plan allow news search + default: false + allow_related_searches: + type: boolean + required: false + label: + en_US: Allow Related Searches + zh_Hans: 支持相关搜索 + pt_BR: Allow Related Searches + help: + en_US: Does your subscription plan allow related searches + zh_Hans: 您的订阅计划是否支持相关搜索 + pt_BR: Does your subscription plan allow related searches + default: false diff --git a/api/core/tools/provider/builtin/bing/tools/bing_web_search.py b/api/core/tools/provider/builtin/bing/tools/bing_web_search.py new file mode 100644 index 0000000000000000000000000000000000000000..0de693698363f871ec2ba04cfb6d2ecfa474f448 --- /dev/null +++ b/api/core/tools/provider/builtin/bing/tools/bing_web_search.py @@ -0,0 +1,237 @@ +from typing import Any, Union +from urllib.parse import quote + +from requests import get + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class BingSearchTool(BuiltinTool): + url: str = "https://api.bing.microsoft.com/v7.0/search" + + def _invoke_bing( + self, + user_id: str, + server_url: str, + subscription_key: str, + query: str, + limit: int, + result_type: str, + market: str, + lang: str, + filters: list[str], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke bing search + """ + market_code = f"{lang}-{market}" + accept_language = f"{lang},{market_code};q=0.9" + headers = {"Ocp-Apim-Subscription-Key": subscription_key, "Accept-Language": accept_language} + + query = quote(query) + server_url = f"{server_url}?q={query}&mkt={market_code}&count={limit}&responseFilter={','.join(filters)}" + response = get(server_url, headers=headers) + + if response.status_code != 200: + raise Exception(f"Error {response.status_code}: {response.text}") + + response = response.json() + search_results = response["webPages"]["value"][:limit] if "webPages" in response else [] + related_searches = response["relatedSearches"]["value"] if "relatedSearches" in response else [] + entities = response["entities"]["value"] if "entities" in response else [] + news = response["news"]["value"] if "news" in response else [] + computation = response["computation"]["value"] if "computation" in response else None + + if result_type == "link": + results = [] + if search_results: + for result in search_results: + url = f": {result['url']}" if "url" in result else "" + results.append(self.create_text_message(text=f"{result['name']}{url}")) + + if entities: + for entity in entities: + url = f": {entity['url']}" if "url" in entity else "" + results.append(self.create_text_message(text=f"{entity.get('name', '')}{url}")) + + if news: + for news_item in news: + url = f": {news_item['url']}" if "url" in news_item else "" + results.append(self.create_text_message(text=f"{news_item.get('name', '')}{url}")) + + if related_searches: + for related in related_searches: + url = f": {related['displayText']}" if "displayText" in related else "" + results.append(self.create_text_message(text=f"{related.get('displayText', '')}{url}")) + + return results + elif result_type == "json": + result = {} + if search_results: + result["organic"] = [ + { + "title": item.get("name", ""), + "snippet": item.get("snippet", ""), + "url": item.get("url", ""), + "siteName": item.get("siteName", ""), + } + for item in search_results + ] + + if computation and "expression" in computation and "value" in computation: + result["computation"] = {"expression": computation["expression"], "value": computation["value"]} + + if entities: + result["entities"] = [ + { + "name": item.get("name", ""), + "url": item.get("url", ""), + "description": item.get("description", ""), + } + for item in entities + ] + + if news: + result["news"] = [{"name": item.get("name", ""), "url": item.get("url", "")} for item in news] + + if related_searches: + result["related searches"] = [ + {"displayText": item.get("displayText", ""), "url": item.get("webSearchUrl", "")} for item in news + ] + + return self.create_json_message(result) + else: + # construct text + text = "" + if search_results: + for i, result in enumerate(search_results): + text += f"{i + 1}: {result.get('name', '')} - {result.get('snippet', '')}\n" + + if computation and "expression" in computation and "value" in computation: + text += "\nComputation:\n" + text += f"{computation['expression']} = {computation['value']}\n" + + if entities: + text += "\nEntities:\n" + for entity in entities: + url = f"- {entity['url']}" if "url" in entity else "" + text += f"{entity.get('name', '')}{url}\n" + + if news: + text += "\nNews:\n" + for news_item in news: + url = f"- {news_item['url']}" if "url" in news_item else "" + text += f"{news_item.get('name', '')}{url}\n" + + if related_searches: + text += "\n\nRelated Searches:\n" + for related in related_searches: + url = f"- {related['webSearchUrl']}" if "webSearchUrl" in related else "" + text += f"{related.get('displayText', '')}{url}\n" + + return self.create_text_message(text=self.summary(user_id=user_id, content=text)) + + def validate_credentials(self, credentials: dict[str, Any], tool_parameters: dict[str, Any]) -> None: + key = credentials.get("subscription_key") + if not key: + raise Exception("subscription_key is required") + + server_url = credentials.get("server_url") + if not server_url: + server_url = self.url + + query = tool_parameters.get("query") + if not query: + raise Exception("query is required") + + limit = min(tool_parameters.get("limit", 5), 10) + result_type = tool_parameters.get("result_type", "text") or "text" + + market = tool_parameters.get("market", "US") + lang = tool_parameters.get("language", "en") + filter = [] + + if credentials.get("allow_entities", False): + filter.append("Entities") + + if credentials.get("allow_computation", False): + filter.append("Computation") + + if credentials.get("allow_news", False): + filter.append("News") + + if credentials.get("allow_related_searches", False): + filter.append("RelatedSearches") + + if credentials.get("allow_web_pages", False): + filter.append("WebPages") + + if not filter: + raise Exception("At least one filter is required") + + self._invoke_bing( + user_id="test", + server_url=server_url, + subscription_key=key, + query=query, + limit=limit, + result_type=result_type, + market=market, + lang=lang, + filters=filter, + ) + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + + key = self.runtime.credentials.get("subscription_key", None) + if not key: + raise Exception("subscription_key is required") + + server_url = self.runtime.credentials.get("server_url", None) + if not server_url: + server_url = self.url + + query = tool_parameters.get("query") + if not query: + raise Exception("query is required") + + limit = min(tool_parameters.get("limit", 5), 10) + result_type = tool_parameters.get("result_type", "text") or "text" + + market = tool_parameters.get("market", "US") + lang = tool_parameters.get("language", "en") + filter = [] + + if tool_parameters.get("enable_computation", False): + filter.append("Computation") + if tool_parameters.get("enable_entities", False): + filter.append("Entities") + if tool_parameters.get("enable_news", False): + filter.append("News") + if tool_parameters.get("enable_related_search", False): + filter.append("RelatedSearches") + if tool_parameters.get("enable_webpages", False): + filter.append("WebPages") + + if not filter: + raise Exception("At least one filter is required") + + return self._invoke_bing( + user_id=user_id, + server_url=server_url, + subscription_key=key, + query=query, + limit=limit, + result_type=result_type, + market=market, + lang=lang, + filters=filter, + ) diff --git a/api/core/tools/provider/builtin/bing/tools/bing_web_search.yaml b/api/core/tools/provider/builtin/bing/tools/bing_web_search.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f5c932c37b3014456693b0258eec966f346044c1 --- /dev/null +++ b/api/core/tools/provider/builtin/bing/tools/bing_web_search.yaml @@ -0,0 +1,589 @@ +identity: + name: bing_web_search + author: Dify + label: + en_US: BingWebSearch + zh_Hans: 必应网页搜索 + pt_BR: BingWebSearch +description: + human: + en_US: A tool for performing a Bing SERP search and extracting snippets and webpages.Input should be a search query. + zh_Hans: 一个用于执行 Bing SERP 搜索并提取片段和网页的工具。输入应该是一个搜索查询。 + pt_BR: A tool for performing a Bing SERP search and extracting snippets and webpages.Input should be a search query. + llm: A tool for performing a Bing SERP search and extracting snippets and webpages.Input should be a search query. +parameters: + - name: query + type: string + required: true + form: llm + label: + en_US: Query string + zh_Hans: 查询语句 + pt_BR: Query string + human_description: + en_US: used for searching + zh_Hans: 用于搜索网页内容 + pt_BR: used for searching + llm_description: key words for searching + - name: enable_computation + type: boolean + required: false + form: form + label: + en_US: Enable computation + zh_Hans: 启用计算 + pt_BR: Enable computation + human_description: + en_US: enable computation + zh_Hans: 启用计算 + pt_BR: enable computation + default: false + - name: enable_entities + type: boolean + required: false + form: form + label: + en_US: Enable entities + zh_Hans: 启用实体搜索 + pt_BR: Enable entities + human_description: + en_US: enable entities + zh_Hans: 启用实体搜索 + pt_BR: enable entities + default: true + - name: enable_news + type: boolean + required: false + form: form + label: + en_US: Enable news + zh_Hans: 启用新闻搜索 + pt_BR: Enable news + human_description: + en_US: enable news + zh_Hans: 启用新闻搜索 + pt_BR: enable news + default: false + - name: enable_related_search + type: boolean + required: false + form: form + label: + en_US: Enable related search + zh_Hans: 启用相关搜索 + pt_BR: Enable related search + human_description: + en_US: enable related search + zh_Hans: 启用相关搜索 + pt_BR: enable related search + default: false + - name: enable_webpages + type: boolean + required: false + form: form + label: + en_US: Enable webpages search + zh_Hans: 启用网页搜索 + pt_BR: Enable webpages search + human_description: + en_US: enable webpages search + zh_Hans: 启用网页搜索 + pt_BR: enable webpages search + default: true + - name: limit + type: number + required: true + form: form + label: + en_US: Limit for results length + zh_Hans: 返回长度限制 + pt_BR: Limit for results length + human_description: + en_US: limit the number of results + zh_Hans: 限制返回结果的数量 + pt_BR: limit the number of results + min: 1 + max: 10 + default: 5 + - name: result_type + type: select + required: true + label: + en_US: result type + zh_Hans: 结果类型 + pt_BR: result type + human_description: + en_US: return a list of links, json or texts + zh_Hans: 返回一个列表,内容是链接、json还是纯文本 + pt_BR: return a list of links, json or texts + default: text + options: + - value: link + label: + en_US: Link + zh_Hans: 链接 + pt_BR: Link + - value: json + label: + en_US: JSON + zh_Hans: JSON + pt_BR: JSON + - value: text + label: + en_US: Text + zh_Hans: 文本 + pt_BR: Text + form: form + - name: market + type: select + label: + en_US: Market + zh_Hans: 市场 + pt_BR: Market + human_description: + en_US: market takes responsibility for the region + zh_Hans: 市场决定了搜索结果的地区 + pt_BR: market takes responsibility for the region + required: false + form: form + default: US + options: + - value: AR + label: + en_US: Argentina + zh_Hans: 阿根廷 + pt_BR: Argentina + - value: AU + label: + en_US: Australia + zh_Hans: 澳大利亚 + pt_BR: Australia + - value: AT + label: + en_US: Austria + zh_Hans: 奥地利 + pt_BR: Austria + - value: BE + label: + en_US: Belgium + zh_Hans: 比利时 + pt_BR: Belgium + - value: BR + label: + en_US: Brazil + zh_Hans: 巴西 + pt_BR: Brazil + - value: CA + label: + en_US: Canada + zh_Hans: 加拿大 + pt_BR: Canada + - value: CL + label: + en_US: Chile + zh_Hans: 智利 + pt_BR: Chile + - value: CO + label: + en_US: Colombia + zh_Hans: 哥伦比亚 + pt_BR: Colombia + - value: CN + label: + en_US: China + zh_Hans: 中国 + pt_BR: China + - value: CZ + label: + en_US: Czech Republic + zh_Hans: 捷克共和国 + pt_BR: Czech Republic + - value: DK + label: + en_US: Denmark + zh_Hans: 丹麦 + pt_BR: Denmark + - value: FI + label: + en_US: Finland + zh_Hans: 芬兰 + pt_BR: Finland + - value: FR + label: + en_US: France + zh_Hans: 法国 + pt_BR: France + - value: DE + label: + en_US: Germany + zh_Hans: 德国 + pt_BR: Germany + - value: HK + label: + en_US: Hong Kong + zh_Hans: 香港 + pt_BR: Hong Kong + - value: IN + label: + en_US: India + zh_Hans: 印度 + pt_BR: India + - value: ID + label: + en_US: Indonesia + zh_Hans: 印度尼西亚 + pt_BR: Indonesia + - value: IT + label: + en_US: Italy + zh_Hans: 意大利 + pt_BR: Italy + - value: JP + label: + en_US: Japan + zh_Hans: 日本 + pt_BR: Japan + - value: KR + label: + en_US: Korea + zh_Hans: 韩国 + pt_BR: Korea + - value: MY + label: + en_US: Malaysia + zh_Hans: 马来西亚 + pt_BR: Malaysia + - value: MX + label: + en_US: Mexico + zh_Hans: 墨西哥 + pt_BR: Mexico + - value: NL + label: + en_US: Netherlands + zh_Hans: 荷兰 + pt_BR: Netherlands + - value: NZ + label: + en_US: New Zealand + zh_Hans: 新西兰 + pt_BR: New Zealand + - value: 'NO' + label: + en_US: Norway + zh_Hans: 挪威 + pt_BR: Norway + - value: PH + label: + en_US: Philippines + zh_Hans: 菲律宾 + pt_BR: Philippines + - value: PL + label: + en_US: Poland + zh_Hans: 波兰 + pt_BR: Poland + - value: PT + label: + en_US: Portugal + zh_Hans: 葡萄牙 + pt_BR: Portugal + - value: RU + label: + en_US: Russia + zh_Hans: 俄罗斯 + pt_BR: Russia + - value: SA + label: + en_US: Saudi Arabia + zh_Hans: 沙特阿拉伯 + pt_BR: Saudi Arabia + - value: SG + label: + en_US: Singapore + zh_Hans: 新加坡 + pt_BR: Singapore + - value: ZA + label: + en_US: South Africa + zh_Hans: 南非 + pt_BR: South Africa + - value: ES + label: + en_US: Spain + zh_Hans: 西班牙 + pt_BR: Spain + - value: SE + label: + en_US: Sweden + zh_Hans: 瑞典 + pt_BR: Sweden + - value: CH + label: + en_US: Switzerland + zh_Hans: 瑞士 + pt_BR: Switzerland + - value: TW + label: + en_US: Taiwan + zh_Hans: 台湾 + pt_BR: Taiwan + - value: TH + label: + en_US: Thailand + zh_Hans: 泰国 + pt_BR: Thailand + - value: TR + label: + en_US: Turkey + zh_Hans: 土耳其 + pt_BR: Turkey + - value: GB + label: + en_US: United Kingdom + zh_Hans: 英国 + pt_BR: United Kingdom + - value: US + label: + en_US: United States + zh_Hans: 美国 + pt_BR: United States + - name: language + type: select + label: + en_US: Language + zh_Hans: 语言 + pt_BR: Language + human_description: + en_US: language takes responsibility for the language of the search result + zh_Hans: 语言决定了搜索结果的语言 + pt_BR: language takes responsibility for the language of the search result + required: false + default: en + form: form + options: + - value: ar + label: + en_US: Arabic + zh_Hans: 阿拉伯语 + pt_BR: Arabic + - value: bg + label: + en_US: Bulgarian + zh_Hans: 保加利亚语 + pt_BR: Bulgarian + - value: ca + label: + en_US: Catalan + zh_Hans: 加泰罗尼亚语 + pt_BR: Catalan + - value: zh-hans + label: + en_US: Chinese (Simplified) + zh_Hans: 中文(简体) + pt_BR: Chinese (Simplified) + - value: zh-hant + label: + en_US: Chinese (Traditional) + zh_Hans: 中文(繁体) + pt_BR: Chinese (Traditional) + - value: cs + label: + en_US: Czech + zh_Hans: 捷克语 + pt_BR: Czech + - value: da + label: + en_US: Danish + zh_Hans: 丹麦语 + pt_BR: Danish + - value: nl + label: + en_US: Dutch + zh_Hans: 荷兰语 + pt_BR: Dutch + - value: en + label: + en_US: English + zh_Hans: 英语 + pt_BR: English + - value: et + label: + en_US: Estonian + zh_Hans: 爱沙尼亚语 + pt_BR: Estonian + - value: fi + label: + en_US: Finnish + zh_Hans: 芬兰语 + pt_BR: Finnish + - value: fr + label: + en_US: French + zh_Hans: 法语 + pt_BR: French + - value: de + label: + en_US: German + zh_Hans: 德语 + pt_BR: German + - value: el + label: + en_US: Greek + zh_Hans: 希腊语 + pt_BR: Greek + - value: he + label: + en_US: Hebrew + zh_Hans: 希伯来语 + pt_BR: Hebrew + - value: hi + label: + en_US: Hindi + zh_Hans: 印地语 + pt_BR: Hindi + - value: hu + label: + en_US: Hungarian + zh_Hans: 匈牙利语 + pt_BR: Hungarian + - value: id + label: + en_US: Indonesian + zh_Hans: 印尼语 + pt_BR: Indonesian + - value: it + label: + en_US: Italian + zh_Hans: 意大利语 + pt_BR: Italian + - value: jp + label: + en_US: Japanese + zh_Hans: 日语 + pt_BR: Japanese + - value: kn + label: + en_US: Kannada + zh_Hans: 卡纳达语 + pt_BR: Kannada + - value: ko + label: + en_US: Korean + zh_Hans: 韩语 + pt_BR: Korean + - value: lv + label: + en_US: Latvian + zh_Hans: 拉脱维亚语 + pt_BR: Latvian + - value: lt + label: + en_US: Lithuanian + zh_Hans: 立陶宛语 + pt_BR: Lithuanian + - value: ms + label: + en_US: Malay + zh_Hans: 马来语 + pt_BR: Malay + - value: ml + label: + en_US: Malayalam + zh_Hans: 马拉雅拉姆语 + pt_BR: Malayalam + - value: mr + label: + en_US: Marathi + zh_Hans: 马拉地语 + pt_BR: Marathi + - value: nb + label: + en_US: Norwegian + zh_Hans: 挪威语 + pt_BR: Norwegian + - value: pl + label: + en_US: Polish + zh_Hans: 波兰语 + pt_BR: Polish + - value: pt-br + label: + en_US: Portuguese (Brazil) + zh_Hans: 葡萄牙语(巴西) + pt_BR: Portuguese (Brazil) + - value: pt-pt + label: + en_US: Portuguese (Portugal) + zh_Hans: 葡萄牙语(葡萄牙) + pt_BR: Portuguese (Portugal) + - value: pa + label: + en_US: Punjabi + zh_Hans: 旁遮普语 + pt_BR: Punjabi + - value: ro + label: + en_US: Romanian + zh_Hans: 罗马尼亚语 + pt_BR: Romanian + - value: ru + label: + en_US: Russian + zh_Hans: 俄语 + pt_BR: Russian + - value: sr + label: + en_US: Serbian + zh_Hans: 塞尔维亚语 + pt_BR: Serbian + - value: sk + label: + en_US: Slovak + zh_Hans: 斯洛伐克语 + pt_BR: Slovak + - value: sl + label: + en_US: Slovenian + zh_Hans: 斯洛文尼亚语 + pt_BR: Slovenian + - value: es + label: + en_US: Spanish + zh_Hans: 西班牙语 + pt_BR: Spanish + - value: sv + label: + en_US: Swedish + zh_Hans: 瑞典语 + pt_BR: Swedish + - value: ta + label: + en_US: Tamil + zh_Hans: 泰米尔语 + pt_BR: Tamil + - value: te + label: + en_US: Telugu + zh_Hans: 泰卢固语 + pt_BR: Telugu + - value: th + label: + en_US: Thai + zh_Hans: 泰语 + pt_BR: Thai + - value: tr + label: + en_US: Turkish + zh_Hans: 土耳其语 + pt_BR: Turkish + - value: uk + label: + en_US: Ukrainian + zh_Hans: 乌克兰语 + pt_BR: Ukrainian + - value: vi + label: + en_US: Vietnamese + zh_Hans: 越南语 + pt_BR: Vietnamese diff --git a/api/core/tools/provider/builtin/brave/_assets/icon.svg b/api/core/tools/provider/builtin/brave/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..d059f7c5161e988c236d0b828e276a455c0af527 --- /dev/null +++ b/api/core/tools/provider/builtin/brave/_assets/icon.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/brave/brave.py b/api/core/tools/provider/builtin/brave/brave.py new file mode 100644 index 0000000000000000000000000000000000000000..c24ee67334083b3487161dd341934cb68f8e04d4 --- /dev/null +++ b/api/core/tools/provider/builtin/brave/brave.py @@ -0,0 +1,22 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.brave.tools.brave_search import BraveSearchTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class BraveProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + BraveSearchTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "query": "Sachin Tendulkar", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/brave/brave.yaml b/api/core/tools/provider/builtin/brave/brave.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2b0dcc0188caf8b23b74979baf4edc2aa95ba4ad --- /dev/null +++ b/api/core/tools/provider/builtin/brave/brave.yaml @@ -0,0 +1,39 @@ +identity: + author: Yash Parmar + name: brave + label: + en_US: Brave + zh_Hans: Brave + pt_BR: Brave + description: + en_US: Brave + zh_Hans: Brave + pt_BR: Brave + icon: icon.svg + tags: + - search +credentials_for_provider: + brave_search_api_key: + type: secret-input + required: true + label: + en_US: Brave Search API key + zh_Hans: Brave Search API key + pt_BR: Brave Search API key + placeholder: + en_US: Please input your Brave Search API key + zh_Hans: 请输入你的 Brave Search API key + pt_BR: Please input your Brave Search API key + help: + en_US: Get your Brave Search API key from Brave + zh_Hans: 从 Brave 获取您的 Brave Search API key + pt_BR: Get your Brave Search API key from Brave + url: https://brave.com/search/api/ + base_url: + type: text-input + required: false + label: + en_US: Brave server's Base URL + zh_Hans: Brave服务器的API URL + placeholder: + en_US: https://api.search.brave.com/res/v1/web/search diff --git a/api/core/tools/provider/builtin/brave/tools/brave_search.py b/api/core/tools/provider/builtin/brave/tools/brave_search.py new file mode 100644 index 0000000000000000000000000000000000000000..c34362ae52ecac95d77bbcd018520917062fee1b --- /dev/null +++ b/api/core/tools/provider/builtin/brave/tools/brave_search.py @@ -0,0 +1,138 @@ +import json +from typing import Any, Optional + +import requests +from pydantic import BaseModel, Field + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +BRAVE_BASE_URL = "https://api.search.brave.com/res/v1/web/search" + + +class BraveSearchWrapper(BaseModel): + """Wrapper around the Brave search engine.""" + + api_key: str + """The API key to use for the Brave search engine.""" + search_kwargs: dict = Field(default_factory=dict) + """Additional keyword arguments to pass to the search request.""" + base_url: str = BRAVE_BASE_URL + """The base URL for the Brave search engine.""" + ensure_ascii: bool = True + """Ensure the JSON output is ASCII encoded.""" + + def run(self, query: str) -> str: + """Query the Brave search engine and return the results as a JSON string. + + Args: + query: The query to search for. + + Returns: The results as a JSON string. + + """ + web_search_results = self._search_request(query=query) + final_results = [ + { + "title": item.get("title"), + "link": item.get("url"), + "snippet": item.get("description"), + } + for item in web_search_results + ] + return json.dumps(final_results, ensure_ascii=self.ensure_ascii) + + def _search_request(self, query: str) -> list[dict]: + headers = { + "X-Subscription-Token": self.api_key, + "Accept": "application/json", + } + req = requests.PreparedRequest() + params = {**self.search_kwargs, **{"q": query}} + req.prepare_url(self.base_url, params) + if req.url is None: + raise ValueError("prepared url is None, this should not happen") + + response = requests.get(req.url, headers=headers) + if not response.ok: + raise Exception(f"HTTP error {response.status_code}") + + return response.json().get("web", {}).get("results", []) + + +class BraveSearch(BaseModel): + """Tool that queries the BraveSearch.""" + + name: str = "brave_search" + description: str = ( + "a search engine. " + "useful for when you need to answer questions about current events." + " input should be a search query." + ) + search_wrapper: BraveSearchWrapper + + @classmethod + def from_api_key( + cls, api_key: str, base_url: str, search_kwargs: Optional[dict] = None, ensure_ascii: bool = True, **kwargs: Any + ) -> "BraveSearch": + """Create a tool from an api key. + + Args: + api_key: The api key to use. + search_kwargs: Any additional kwargs to pass to the search wrapper. + **kwargs: Any additional kwargs to pass to the tool. + + Returns: + A tool. + """ + wrapper = BraveSearchWrapper( + api_key=api_key, base_url=base_url, search_kwargs=search_kwargs or {}, ensure_ascii=ensure_ascii + ) + return cls(search_wrapper=wrapper, **kwargs) + + def _run( + self, + query: str, + ) -> str: + """Use the tool.""" + return self.search_wrapper.run(query) + + +class BraveSearchTool(BuiltinTool): + """ + Tool for performing a search using Brave search engine. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage | list[ToolInvokeMessage]: + """ + Invoke the Brave search tool. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (dict[str, Any]): The parameters for the tool invocation. + + Returns: + ToolInvokeMessage | list[ToolInvokeMessage]: The result of the tool invocation. + """ + query = tool_parameters.get("query", "") + count = tool_parameters.get("count", 3) + api_key = self.runtime.credentials["brave_search_api_key"] + base_url = self.runtime.credentials.get("base_url", BRAVE_BASE_URL) + ensure_ascii = tool_parameters.get("ensure_ascii", True) + + if len(base_url) == 0: + base_url = BRAVE_BASE_URL + + if not query: + return self.create_text_message("Please input query") + + tool = BraveSearch.from_api_key( + api_key=api_key, base_url=base_url, search_kwargs={"count": count}, ensure_ascii=ensure_ascii + ) + + results = tool._run(query) + + if not results: + return self.create_text_message(f"No results found for '{query}' in Tavily") + else: + return self.create_text_message(text=results) diff --git a/api/core/tools/provider/builtin/brave/tools/brave_search.yaml b/api/core/tools/provider/builtin/brave/tools/brave_search.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5222a375f84ceed0ded52157e9b05351d84e3350 --- /dev/null +++ b/api/core/tools/provider/builtin/brave/tools/brave_search.yaml @@ -0,0 +1,53 @@ +identity: + name: brave_search + author: Yash Parmar + label: + en_US: BraveSearch + zh_Hans: BraveSearch + pt_BR: BraveSearch +description: + human: + en_US: BraveSearch is a privacy-focused search engine that leverages its own index to deliver unbiased, independent, and fast search results. It's designed to respect user privacy by not tracking searches or personal information, making it a secure choice for those concerned about online privacy. + zh_Hans: BraveSearch 是一个注重隐私的搜索引擎,它利用自己的索引来提供公正、独立和快速的搜索结果。它旨在通过不跟踪搜索或个人信息来尊重用户隐私,为那些关注在线隐私的用户提供了一个安全的选择。 + pt_BR: BraveSearch é um mecanismo de busca focado na privacidade que utiliza seu próprio índice para entregar resultados de busca imparciais, independentes e rápidos. Ele é projetado para respeitar a privacidade do usuário, não rastreando buscas ou informações pessoais, tornando-se uma escolha segura para aqueles preocupados com a privacidade online. + llm: BraveSearch is a privacy-centric search engine utilizing its unique index to offer unbiased, independent, and swift search results. It aims to protect user privacy by avoiding the tracking of search activities or personal data, presenting a secure option for users mindful of their online privacy. +parameters: + - name: query + type: string + required: true + label: + en_US: Query string + zh_Hans: 查询语句 + pt_BR: Query string + human_description: + en_US: The text input used for initiating searches on the web, focusing on delivering relevant and accurate results without compromising user privacy. + zh_Hans: 用于在网上启动搜索的文本输入,专注于提供相关且准确的结果,同时不妨碍用户隐私。 + pt_BR: A entrada de texto usada para iniciar pesquisas na web, focada em entregar resultados relevantes e precisos sem comprometer a privacidade do usuário. + llm_description: Keywords or phrases entered to perform searches, aimed at providing relevant and precise results while ensuring the privacy of the user is maintained. + form: llm + - name: count + type: number + required: false + default: 3 + label: + en_US: Result count + zh_Hans: 结果数量 + pt_BR: Contagem de resultados + human_description: + en_US: The number of search results to return, allowing users to control the breadth of their search output. + zh_Hans: 要返回的搜索结果数量,允许用户控制他们搜索输出的广度。 + pt_BR: O número de resultados de pesquisa a serem retornados, permitindo que os usuários controlem a amplitude de sua saída de pesquisa. + llm_description: Specifies the amount of search results to be displayed, offering users the ability to adjust the scope of their search findings. + form: llm + - name: ensure_ascii + type: boolean + default: true + label: + en_US: Ensure ASCII + zh_Hans: 确保 ASCII + pt_BR: Ensure ASCII + human_description: + en_US: Ensure the JSON output is ASCII encoded + zh_Hans: 确保输出的 JSON 是 ASCII 编码 + pt_BR: Ensure the JSON output is ASCII encoded + form: form diff --git a/api/core/tools/provider/builtin/chart/_assets/icon.png b/api/core/tools/provider/builtin/chart/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..878e56a0512c31735cc94480cef9b8fa5dfcc7fd Binary files /dev/null and b/api/core/tools/provider/builtin/chart/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/chart/chart.py b/api/core/tools/provider/builtin/chart/chart.py new file mode 100644 index 0000000000000000000000000000000000000000..8fa647d9ed813842ec80ec5ea1dc1244a52ae102 --- /dev/null +++ b/api/core/tools/provider/builtin/chart/chart.py @@ -0,0 +1,38 @@ +import matplotlib +import matplotlib.pyplot as plt +from matplotlib.font_manager import FontProperties, fontManager + +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +def set_chinese_font(): + to_find_fonts = [ + "PingFang SC", + "SimHei", + "Microsoft YaHei", + "STSong", + "SimSun", + "Arial Unicode MS", + "Noto Sans CJK SC", + "Noto Sans CJK JP", + ] + installed_fonts = frozenset(fontInfo.name for fontInfo in fontManager.ttflist) + for font in to_find_fonts: + if font in installed_fonts: + return FontProperties(font) + + return FontProperties() + + +# use non-interactive backend to prevent `RuntimeError: main thread is not in main loop` +matplotlib.use("Agg") +# use a business theme +plt.style.use("seaborn-v0_8-darkgrid") +plt.rcParams["axes.unicode_minus"] = False +font_properties = set_chinese_font() +plt.rcParams["font.family"] = font_properties.get_name() + + +class ChartProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + pass diff --git a/api/core/tools/provider/builtin/chart/chart.yaml b/api/core/tools/provider/builtin/chart/chart.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ad0d9a6cd688cfd61e9b6a8f242eec68c7030912 --- /dev/null +++ b/api/core/tools/provider/builtin/chart/chart.yaml @@ -0,0 +1,17 @@ +identity: + author: Dify + name: chart + label: + en_US: ChartGenerator + zh_Hans: 图表生成 + pt_BR: Gerador de gráficos + description: + en_US: Chart Generator is a tool for generating statistical charts like bar chart, line chart, pie chart, etc. + zh_Hans: 图表生成是一个用于生成可视化图表的工具,你可以通过它来生成柱状图、折线图、饼图等各类图表 + pt_BR: O Gerador de gráficos é uma ferramenta para gerar gráficos estatísticos como gráfico de barras, gráfico de linhas, gráfico de pizza, etc. + icon: icon.png + tags: + - design + - productivity + - utilities +credentials_for_provider: diff --git a/api/core/tools/provider/builtin/chart/tools/bar.py b/api/core/tools/provider/builtin/chart/tools/bar.py new file mode 100644 index 0000000000000000000000000000000000000000..20ce5e138b5bfeca359d1399d0d6e38fab5ca683 --- /dev/null +++ b/api/core/tools/provider/builtin/chart/tools/bar.py @@ -0,0 +1,50 @@ +import io +from typing import Any, Union + +import matplotlib.pyplot as plt + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class BarChartTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + data = tool_parameters.get("data", "") + if not data: + return self.create_text_message("Please input data") + data = data.split(";") + + # if all data is int, convert to int + if all(i.isdigit() for i in data): + data = [int(i) for i in data] + else: + data = [float(i) for i in data] + + axis = tool_parameters.get("x_axis") or None + if axis: + axis = axis.split(";") + if len(axis) != len(data): + axis = None + + flg, ax = plt.subplots(figsize=(10, 8)) + + if axis: + axis = [label[:10] + "..." if len(label) > 10 else label for label in axis] + ax.set_xticklabels(axis, rotation=45, ha="right") + # ensure all labels, including duplicates, are correctly displayed + ax.bar(range(len(data)), data) + ax.set_xticks(range(len(data))) + else: + ax.bar(range(len(data)), data) + + buf = io.BytesIO() + flg.savefig(buf, format="png") + buf.seek(0) + plt.close(flg) + + return [ + self.create_text_message("the bar chart is saved as an image."), + self.create_blob_message(blob=buf.read(), meta={"mime_type": "image/png"}), + ] diff --git a/api/core/tools/provider/builtin/chart/tools/bar.yaml b/api/core/tools/provider/builtin/chart/tools/bar.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ee7405f6810efa6525148c4262822e099a657e0a --- /dev/null +++ b/api/core/tools/provider/builtin/chart/tools/bar.yaml @@ -0,0 +1,41 @@ +identity: + name: bar_chart + author: Dify + label: + en_US: Bar Chart + zh_Hans: 柱状图 + pt_BR: Gráfico de barras + icon: icon.svg +description: + human: + en_US: Bar chart + zh_Hans: 柱状图 + pt_BR: Gráfico de barras + llm: generate a bar chart with input data +parameters: + - name: data + type: string + required: true + label: + en_US: data + zh_Hans: 数据 + pt_BR: dados + human_description: + en_US: data for generating chart, each number should be separated by ";" + zh_Hans: 用于生成柱状图的数据,每个数字之间用 ";" 分隔 + pt_BR: dados para gerar gráfico de barras, cada número deve ser separado por ";" + llm_description: data for generating bar chart, data should be a string contains a list of numbers like "1;2;3;4;5" + form: llm + - name: x_axis + type: string + required: false + label: + en_US: X Axis + zh_Hans: x 轴 + pt_BR: Eixo X + human_description: + en_US: X axis for chart, each text should be separated by ";" + zh_Hans: 柱状图的 x 轴,每个文本之间用 ";" 分隔 + pt_BR: Eixo X para gráfico de barras, cada texto deve ser separado por ";" + llm_description: x axis for bar chart, x axis should be a string contains a list of texts like "a;b;c;1;2" in order to match the data + form: llm diff --git a/api/core/tools/provider/builtin/chart/tools/line.py b/api/core/tools/provider/builtin/chart/tools/line.py new file mode 100644 index 0000000000000000000000000000000000000000..39e8caac7ef609f8b09a55f8ebe4ddcbe208f676 --- /dev/null +++ b/api/core/tools/provider/builtin/chart/tools/line.py @@ -0,0 +1,50 @@ +import io +from typing import Any, Union + +import matplotlib.pyplot as plt + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class LinearChartTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + data = tool_parameters.get("data", "") + if not data: + return self.create_text_message("Please input data") + data = data.split(";") + + axis = tool_parameters.get("x_axis") or None + if axis: + axis = axis.split(";") + if len(axis) != len(data): + axis = None + + # if all data is int, convert to int + if all(i.isdigit() for i in data): + data = [int(i) for i in data] + else: + data = [float(i) for i in data] + + flg, ax = plt.subplots(figsize=(10, 8)) + + if axis: + axis = [label[:10] + "..." if len(label) > 10 else label for label in axis] + ax.set_xticklabels(axis, rotation=45, ha="right") + ax.plot(axis, data) + else: + ax.plot(data) + + buf = io.BytesIO() + flg.savefig(buf, format="png") + buf.seek(0) + plt.close(flg) + + return [ + self.create_text_message("the linear chart is saved as an image."), + self.create_blob_message(blob=buf.read(), meta={"mime_type": "image/png"}), + ] diff --git a/api/core/tools/provider/builtin/chart/tools/line.yaml b/api/core/tools/provider/builtin/chart/tools/line.yaml new file mode 100644 index 0000000000000000000000000000000000000000..35ebe3b68bddb35561e25de3537a90b68a50eb15 --- /dev/null +++ b/api/core/tools/provider/builtin/chart/tools/line.yaml @@ -0,0 +1,41 @@ +identity: + name: line_chart + author: Dify + label: + en_US: Linear Chart + zh_Hans: 线性图表 + pt_BR: Gráfico linear + icon: icon.svg +description: + human: + en_US: linear chart + zh_Hans: 线性图表 + pt_BR: Gráfico linear + llm: generate a linear chart with input data +parameters: + - name: data + type: string + required: true + label: + en_US: data + zh_Hans: 数据 + pt_BR: dados + human_description: + en_US: data for generating chart, each number should be separated by ";" + zh_Hans: 用于生成线性图表的数据,每个数字之间用 ";" 分隔 + pt_BR: dados para gerar gráfico linear, cada número deve ser separado por ";" + llm_description: data for generating linear chart, data should be a string contains a list of numbers like "1;2;3;4;5" + form: llm + - name: x_axis + type: string + required: false + label: + en_US: X Axis + zh_Hans: x 轴 + pt_BR: Eixo X + human_description: + en_US: X axis for chart, each text should be separated by ";" + zh_Hans: 线性图表的 x 轴,每个文本之间用 ";" 分隔 + pt_BR: Eixo X para gráfico linear, cada texto deve ser separado por ";" + llm_description: x axis for linear chart, x axis should be a string contains a list of texts like "a;b;c;1;2" in order to match the data + form: llm diff --git a/api/core/tools/provider/builtin/chart/tools/pie.py b/api/core/tools/provider/builtin/chart/tools/pie.py new file mode 100644 index 0000000000000000000000000000000000000000..2c3b8a733eac9a25b11c70c1da6ab2a566a1a328 --- /dev/null +++ b/api/core/tools/provider/builtin/chart/tools/pie.py @@ -0,0 +1,48 @@ +import io +from typing import Any, Union + +import matplotlib.pyplot as plt + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class PieChartTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + data = tool_parameters.get("data", "") + if not data: + return self.create_text_message("Please input data") + data = data.split(";") + categories = tool_parameters.get("categories") or None + + # if all data is int, convert to int + if all(i.isdigit() for i in data): + data = [int(i) for i in data] + else: + data = [float(i) for i in data] + + flg, ax = plt.subplots() + + if categories: + categories = categories.split(";") + if len(categories) != len(data): + categories = None + + if categories: + ax.pie(data, labels=categories) + else: + ax.pie(data) + + buf = io.BytesIO() + flg.savefig(buf, format="png") + buf.seek(0) + plt.close(flg) + + return [ + self.create_text_message("the pie chart is saved as an image."), + self.create_blob_message(blob=buf.read(), meta={"mime_type": "image/png"}), + ] diff --git a/api/core/tools/provider/builtin/chart/tools/pie.yaml b/api/core/tools/provider/builtin/chart/tools/pie.yaml new file mode 100644 index 0000000000000000000000000000000000000000..541715cb7d86dd93b7b7a3d0d3cbf0918e5f97ae --- /dev/null +++ b/api/core/tools/provider/builtin/chart/tools/pie.yaml @@ -0,0 +1,41 @@ +identity: + name: pie_chart + author: Dify + label: + en_US: Pie Chart + zh_Hans: 饼图 + pt_BR: Gráfico de pizza + icon: icon.svg +description: + human: + en_US: Pie chart + zh_Hans: 饼图 + pt_BR: Gráfico de pizza + llm: generate a pie chart with input data +parameters: + - name: data + type: string + required: true + label: + en_US: data + zh_Hans: 数据 + pt_BR: dados + human_description: + en_US: data for generating chart, each number should be separated by ";" + zh_Hans: 用于生成饼图的数据,每个数字之间用 ";" 分隔 + pt_BR: dados para gerar gráfico de pizza, cada número deve ser separado por ";" + llm_description: data for generating pie chart, data should be a string contains a list of numbers like "1;2;3;4;5" + form: llm + - name: categories + type: string + required: true + label: + en_US: Categories + zh_Hans: 分类 + pt_BR: Categorias + human_description: + en_US: Categories for chart, each category should be separated by ";" + zh_Hans: 饼图的分类,每个分类之间用 ";" 分隔 + pt_BR: Categorias para gráfico de pizza, cada categoria deve ser separada por ";" + llm_description: categories for pie chart, categories should be a string contains a list of texts like "a;b;c;1;2" in order to match the data, each category should be split by ";" + form: llm diff --git a/api/core/tools/provider/builtin/code/_assets/icon.svg b/api/core/tools/provider/builtin/code/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..b986ed942621c9ff862d1a664e77cf1e429bdf2c --- /dev/null +++ b/api/core/tools/provider/builtin/code/_assets/icon.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/code/code.py b/api/core/tools/provider/builtin/code/code.py new file mode 100644 index 0000000000000000000000000000000000000000..211417c9a431ed7c3d21ac39bf0ac555e67f2086 --- /dev/null +++ b/api/core/tools/provider/builtin/code/code.py @@ -0,0 +1,8 @@ +from typing import Any + +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class CodeToolProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + pass diff --git a/api/core/tools/provider/builtin/code/code.yaml b/api/core/tools/provider/builtin/code/code.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2640a7087ef28bf9499fc1b57ee7f5e18281d3f7 --- /dev/null +++ b/api/core/tools/provider/builtin/code/code.yaml @@ -0,0 +1,15 @@ +identity: + author: Dify + name: code + label: + en_US: Code Interpreter + zh_Hans: 代码解释器 + pt_BR: Interpretador de Código + description: + en_US: Run a piece of code and get the result back. + zh_Hans: 运行一段代码并返回结果。 + pt_BR: Execute um trecho de código e obtenha o resultado de volta. + icon: icon.svg + tags: + - productivity +credentials_for_provider: diff --git a/api/core/tools/provider/builtin/code/tools/simple_code.py b/api/core/tools/provider/builtin/code/tools/simple_code.py new file mode 100644 index 0000000000000000000000000000000000000000..632c9fc7f1451bc311e9383667bef6cc86b9ac7c --- /dev/null +++ b/api/core/tools/provider/builtin/code/tools/simple_code.py @@ -0,0 +1,22 @@ +from typing import Any + +from core.helper.code_executor.code_executor import CodeExecutor, CodeLanguage +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class SimpleCode(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage | list[ToolInvokeMessage]: + """ + invoke simple code + """ + + language = tool_parameters.get("language", CodeLanguage.PYTHON3) + code = tool_parameters.get("code", "") + + if language not in {CodeLanguage.PYTHON3, CodeLanguage.JAVASCRIPT}: + raise ValueError(f"Only python3 and javascript are supported, not {language}") + + result = CodeExecutor.execute_code(language, "", code) + + return self.create_text_message(result) diff --git a/api/core/tools/provider/builtin/code/tools/simple_code.yaml b/api/core/tools/provider/builtin/code/tools/simple_code.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0f51674987d7d4f09483740c4fc8df5534a9c4cd --- /dev/null +++ b/api/core/tools/provider/builtin/code/tools/simple_code.yaml @@ -0,0 +1,51 @@ +identity: + name: simple_code + author: Dify + label: + en_US: Code Interpreter + zh_Hans: 代码解释器 + pt_BR: Interpretador de Código +description: + human: + en_US: Run code and get the result back. When you're using a lower quality model, please make sure there are some tips help LLM to understand how to write the code. + zh_Hans: 运行一段代码并返回结果。当您使用较低质量的模型时,请确保有一些提示帮助LLM理解如何编写代码。 + pt_BR: Execute um trecho de código e obtenha o resultado de volta. quando você estiver usando um modelo de qualidade inferior, certifique-se de que existam algumas dicas para ajudar o LLM a entender como escrever o código. + llm: A tool for running code and getting the result back. Only native packages are allowed, network/IO operations are disabled. and you must use print() or console.log() to output the result or result will be empty. +parameters: + - name: language + type: string + required: true + label: + en_US: Language + zh_Hans: 语言 + pt_BR: Idioma + human_description: + en_US: The programming language of the code + zh_Hans: 代码的编程语言 + pt_BR: A linguagem de programação do código + llm_description: language of the code, only "python3" and "javascript" are supported + form: llm + options: + - value: python3 + label: + en_US: Python3 + zh_Hans: Python3 + pt_BR: Python3 + - value: javascript + label: + en_US: JavaScript + zh_Hans: JavaScript + pt_BR: JavaScript + - name: code + type: string + required: true + label: + en_US: Code + zh_Hans: 代码 + pt_BR: Código + human_description: + en_US: The code to be executed + zh_Hans: 要执行的代码 + pt_BR: O código a ser executado + llm_description: code to be executed, only native packages are allowed, network/IO operations are disabled. + form: llm diff --git a/api/core/tools/provider/builtin/cogview/__init__.py b/api/core/tools/provider/builtin/cogview/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/tools/provider/builtin/cogview/_assets/icon.png b/api/core/tools/provider/builtin/cogview/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..f0c1c24a02fc838655e47d58dd00a25a41620e78 Binary files /dev/null and b/api/core/tools/provider/builtin/cogview/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/cogview/cogview.py b/api/core/tools/provider/builtin/cogview/cogview.py new file mode 100644 index 0000000000000000000000000000000000000000..6941ce864956937ded6b815bf5840b29a2123510 --- /dev/null +++ b/api/core/tools/provider/builtin/cogview/cogview.py @@ -0,0 +1,28 @@ +"""Provide the input parameters type for the cogview provider class""" + +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.cogview.tools.cogview3 import CogView3Tool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class COGVIEWProvider(BuiltinToolProviderController): + """cogview provider""" + + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + CogView3Tool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "prompt": "一个城市在水晶瓶中欢快生活的场景,水彩画风格,展现出微观与珠宝般的美丽。", + "size": "square", + "n": 1, + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) from e diff --git a/api/core/tools/provider/builtin/cogview/cogview.yaml b/api/core/tools/provider/builtin/cogview/cogview.yaml new file mode 100644 index 0000000000000000000000000000000000000000..374b0e98d9122ce40a59e6c3cf09404a7cd5fae6 --- /dev/null +++ b/api/core/tools/provider/builtin/cogview/cogview.yaml @@ -0,0 +1,61 @@ +identity: + author: Waffle + name: cogview + label: + en_US: CogView + zh_Hans: CogView 绘画 + pt_BR: CogView + description: + en_US: CogView art + zh_Hans: CogView 绘画 + pt_BR: CogView art + icon: icon.png + tags: + - image + - productivity +credentials_for_provider: + zhipuai_api_key: + type: secret-input + required: true + label: + en_US: ZhipuAI API key + zh_Hans: ZhipuAI API key + pt_BR: ZhipuAI API key + help: + en_US: Please input your ZhipuAI API key + zh_Hans: 请输入你的 ZhipuAI API key + pt_BR: Please input your ZhipuAI API key + placeholder: + en_US: Please input your ZhipuAI API key + zh_Hans: 请输入你的 ZhipuAI API key + pt_BR: Please input your ZhipuAI API key + zhipuai_organizaion_id: + type: text-input + required: false + label: + en_US: ZhipuAI organization ID + zh_Hans: ZhipuAI organization ID + pt_BR: ZhipuAI organization ID + help: + en_US: Please input your ZhipuAI organization ID + zh_Hans: 请输入你的 ZhipuAI organization ID + pt_BR: Please input your ZhipuAI organization ID + placeholder: + en_US: Please input your ZhipuAI organization ID + zh_Hans: 请输入你的 ZhipuAI organization ID + pt_BR: Please input your ZhipuAI organization ID + zhipuai_base_url: + type: text-input + required: false + label: + en_US: ZhipuAI base URL + zh_Hans: ZhipuAI base URL + pt_BR: ZhipuAI base URL + help: + en_US: Please input your ZhipuAI base URL + zh_Hans: 请输入你的 ZhipuAI base URL + pt_BR: Please input your ZhipuAI base URL + placeholder: + en_US: Please input your ZhipuAI base URL + zh_Hans: 请输入你的 ZhipuAI base URL + pt_BR: Please input your ZhipuAI base URL diff --git a/api/core/tools/provider/builtin/cogview/tools/cogvideo.py b/api/core/tools/provider/builtin/cogview/tools/cogvideo.py new file mode 100644 index 0000000000000000000000000000000000000000..a60062ca66abbfc13b4c2400d9f8d885b1fcf4ac --- /dev/null +++ b/api/core/tools/provider/builtin/cogview/tools/cogvideo.py @@ -0,0 +1,24 @@ +from typing import Any, Union + +from zhipuai import ZhipuAI # type: ignore + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class CogVideoTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + client = ZhipuAI( + base_url=self.runtime.credentials["zhipuai_base_url"], + api_key=self.runtime.credentials["zhipuai_api_key"], + ) + if not tool_parameters.get("prompt") and not tool_parameters.get("image_url"): + return self.create_text_message("require at least one of prompt and image_url") + + response = client.videos.generations( + model="cogvideox", prompt=tool_parameters.get("prompt"), image_url=tool_parameters.get("image_url") + ) + + return self.create_json_message(response.dict()) diff --git a/api/core/tools/provider/builtin/cogview/tools/cogvideo.yaml b/api/core/tools/provider/builtin/cogview/tools/cogvideo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3df0cfcea938fad721a8be8b037716853df272c4 --- /dev/null +++ b/api/core/tools/provider/builtin/cogview/tools/cogvideo.yaml @@ -0,0 +1,32 @@ +identity: + name: cogvideo + author: hjlarry + label: + en_US: CogVideo + zh_Hans: CogVideo 视频生成 +description: + human: + en_US: Use the CogVideox model provided by ZhipuAI to generate videos based on user prompts and images. + zh_Hans: 使用智谱cogvideox模型,根据用户输入的提示词和图片,生成视频。 + llm: A tool for generating videos. The input is user's prompt or image url or both of them, the output is a task id. You can use another tool with this task id to check the status and get the video. +parameters: + - name: prompt + type: string + label: + en_US: prompt + zh_Hans: 提示词 + human_description: + en_US: The prompt text used to generate video. + zh_Hans: 用于生成视频的提示词。 + llm_description: The prompt text used to generate video. Optional. + form: llm + - name: image_url + type: string + label: + en_US: image url + zh_Hans: 图片链接 + human_description: + en_US: The image url used to generate video. + zh_Hans: 输入一个图片链接,生成的视频将基于该图片和提示词。 + llm_description: The image url used to generate video. Optional. + form: llm diff --git a/api/core/tools/provider/builtin/cogview/tools/cogvideo_job.py b/api/core/tools/provider/builtin/cogview/tools/cogvideo_job.py new file mode 100644 index 0000000000000000000000000000000000000000..3e24b74d2598a7e35f5c0586073bf07a64f7bc45 --- /dev/null +++ b/api/core/tools/provider/builtin/cogview/tools/cogvideo_job.py @@ -0,0 +1,30 @@ +from typing import Any, Union + +import httpx +from zhipuai import ZhipuAI # type: ignore + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class CogVideoJobTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + client = ZhipuAI( + api_key=self.runtime.credentials["zhipuai_api_key"], + base_url=self.runtime.credentials["zhipuai_base_url"], + ) + + response = client.videos.retrieve_videos_result(id=tool_parameters.get("id")) + result = [self.create_json_message(response.dict())] + if response.task_status == "SUCCESS": + for item in response.video_result: + video_cover_image = self.create_image_message(item.cover_image_url) + result.append(video_cover_image) + video = self.create_blob_message( + blob=httpx.get(item.url).content, meta={"mime_type": "video/mp4"}, save_as=self.VariableKey.VIDEO + ) + result.append(video) + + return result diff --git a/api/core/tools/provider/builtin/cogview/tools/cogvideo_job.yaml b/api/core/tools/provider/builtin/cogview/tools/cogvideo_job.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fb2eb3ab130b81203915d57deff0c44d81ec6cdf --- /dev/null +++ b/api/core/tools/provider/builtin/cogview/tools/cogvideo_job.yaml @@ -0,0 +1,21 @@ +identity: + name: cogvideo_job + author: hjlarry + label: + en_US: CogVideo Result + zh_Hans: CogVideo 结果获取 +description: + human: + en_US: Get the result of CogVideo tool generation. + zh_Hans: 根据 CogVideo 工具返回的 id 获取视频生成结果。 + llm: Get the result of CogVideo tool generation. The input is the id which is returned by the CogVideo tool. The output is the url of video and video cover image. +parameters: + - name: id + type: string + label: + en_US: id + human_description: + en_US: The id returned by the CogVideo. + zh_Hans: CogVideo 工具返回的 id。 + llm_description: The id returned by the cogvideo. + form: llm diff --git a/api/core/tools/provider/builtin/cogview/tools/cogview3.py b/api/core/tools/provider/builtin/cogview/tools/cogview3.py new file mode 100644 index 0000000000000000000000000000000000000000..9aa781709a726c5f91d8033754057b60681be607 --- /dev/null +++ b/api/core/tools/provider/builtin/cogview/tools/cogview3.py @@ -0,0 +1,93 @@ +import random +from typing import Any, Union + +from zhipuai import ZhipuAI # type: ignore + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class CogView3Tool(BuiltinTool): + """CogView3 Tool""" + + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Invoke CogView3 tool + """ + client = ZhipuAI( + base_url=self.runtime.credentials["zhipuai_base_url"], + api_key=self.runtime.credentials["zhipuai_api_key"], + ) + size_mapping = { + "square": "1024x1024", + "vertical_768": "768x1344", + "vertical_864": "864x1152", + "horizontal_1344": "1344x768", + "horizontal_1152": "1152x864", + "widescreen_1440": "1440x720", + "tallscreen_720": "720x1440", + } + # prompt + prompt = tool_parameters.get("prompt", "") + if not prompt: + return self.create_text_message("Please input prompt") + # get size key + size_key = tool_parameters.get("size", "square") + # cogview-3-plus get size + if size_key != "cogview_3": + size = size_mapping[size_key] + # get n + n = tool_parameters.get("n", 1) + # get quality + quality = tool_parameters.get("quality", "standard") + if quality not in {"standard", "hd"}: + return self.create_text_message("Invalid quality") + # get style + style = tool_parameters.get("style", "vivid") + if style not in {"natural", "vivid"}: + return self.create_text_message("Invalid style") + # set extra body + seed_id = tool_parameters.get("seed_id", self._generate_random_id(8)) + extra_body = {"seed": seed_id} + # cogview-3-plus + if size_key != "cogview_3": + response = client.images.generations( + prompt=prompt, + model="cogview-3-plus", + size=size, + n=n, + extra_body=extra_body, + style=style, + quality=quality, + response_format="b64_json", + ) + # cogview-3 + else: + response = client.images.generations( + prompt=prompt, + model="cogview-3", + n=n, + extra_body=extra_body, + style=style, + quality=quality, + response_format="b64_json", + ) + result = [] + for image in response.data: + result.append(self.create_image_message(image=image.url)) + result.append( + self.create_json_message( + { + "url": image.url, + } + ) + ) + return result + + @staticmethod + def _generate_random_id(length=8): + characters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789" + random_id = "".join(random.choices(characters, k=length)) + return random_id diff --git a/api/core/tools/provider/builtin/cogview/tools/cogview3.yaml b/api/core/tools/provider/builtin/cogview/tools/cogview3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9ab5c2729bf7a945f9b5de8d605508231ad7c689 --- /dev/null +++ b/api/core/tools/provider/builtin/cogview/tools/cogview3.yaml @@ -0,0 +1,148 @@ +identity: + name: cogview3 + author: Waffle + label: + en_US: CogView 3 + zh_Hans: CogView 3 绘画 + pt_BR: CogView 3 + description: + en_US: CogView 3 is a powerful drawing tool that can draw the image you want based on your prompt + zh_Hans: CogView 3 是一个强大的绘画工具,它可以根据您的提示词绘制出您想要的图像 + pt_BR: CogView 3 is a powerful drawing tool that can draw the image you want based on your prompt +description: + human: + en_US: CogView 3 is a text to image tool + zh_Hans: CogView 3 是一个文本到图像的工具 + pt_BR: CogView 3 is a text to image tool + llm: CogView 3 is a tool used to generate images from text +parameters: + - name: prompt + type: string + required: true + label: + en_US: Prompt + zh_Hans: 提示词 + pt_BR: Prompt + human_description: + en_US: Image prompt, you can check the official documentation of CogView 3 + zh_Hans: 图像提示词,您可以查看 CogView 3 的官方文档 + pt_BR: Image prompt, you can check the official documentation of CogView 3 + llm_description: Image prompt of CogView 3, you should describe the image you want to generate as a list of words as possible as detailed + form: llm + - name: size + type: select + required: true + human_description: + en_US: selecting the image size + zh_Hans: 选择图像大小 + pt_BR: selecting the image size + label: + en_US: Image size + zh_Hans: 图像大小 + pt_BR: Image size + form: form + options: + - value: cogview_3 + label: + en_US: Square_cogview_3(1024x1024) + zh_Hans: 方_cogview_3(1024x1024) + pt_BR: Square_cogview_3(1024x1024) + - value: square + label: + en_US: Square(1024x1024) + zh_Hans: 方(1024x1024) + pt_BR: Square(1024x1024) + - value: vertical_768 + label: + en_US: Vertical(768x1344) + zh_Hans: 竖屏(768x1344) + pt_BR: Vertical(768x1344) + - value: vertical_864 + label: + en_US: Vertical(864x1152) + zh_Hans: 竖屏(864x1152) + pt_BR: Vertical(864x1152) + - value: horizontal_1344 + label: + en_US: Horizontal(1344x768) + zh_Hans: 横屏(1344x768) + pt_BR: Horizontal(1344x768) + - value: horizontal_1152 + label: + en_US: Horizontal(1152x864) + zh_Hans: 横屏(1152x864) + pt_BR: Horizontal(1152x864) + - value: widescreen_1440 + label: + en_US: Widescreen(1440x720) + zh_Hans: 宽屏(1440x720) + pt_BR: Widescreen(1440x720) + - value: tallscreen_720 + label: + en_US: Tallscreen(720x1440) + zh_Hans: 高屏(720x1440) + pt_BR: Tallscreen(720x1440) + default: square + - name: n + type: number + required: true + human_description: + en_US: selecting the number of images + zh_Hans: 选择图像数量 + pt_BR: selecting the number of images + label: + en_US: Number of images + zh_Hans: 图像数量 + pt_BR: Number of images + form: form + min: 1 + max: 1 + default: 1 + - name: quality + type: select + required: true + human_description: + en_US: selecting the image quality + zh_Hans: 选择图像质量 + pt_BR: selecting the image quality + label: + en_US: Image quality + zh_Hans: 图像质量 + pt_BR: Image quality + form: form + options: + - value: standard + label: + en_US: Standard + zh_Hans: 标准 + pt_BR: Standard + - value: hd + label: + en_US: HD + zh_Hans: 高清 + pt_BR: HD + default: standard + - name: style + type: select + required: true + human_description: + en_US: selecting the image style + zh_Hans: 选择图像风格 + pt_BR: selecting the image style + label: + en_US: Image style + zh_Hans: 图像风格 + pt_BR: Image style + form: form + options: + - value: vivid + label: + en_US: Vivid + zh_Hans: 生动 + pt_BR: Vivid + - value: natural + label: + en_US: Natural + zh_Hans: 自然 + pt_BR: Natural + default: vivid diff --git a/api/core/tools/provider/builtin/comfyui/_assets/icon.png b/api/core/tools/provider/builtin/comfyui/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..5a98fa13c528e84bf42625c66ae27b49c0fd506c --- /dev/null +++ b/api/core/tools/provider/builtin/comfyui/_assets/icon.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6f55a6a4854e3ffce2698d0913f37fe2f0d80de71a3e2df81cba7ccd920c4b9f +size 213986 diff --git a/api/core/tools/provider/builtin/comfyui/comfyui.py b/api/core/tools/provider/builtin/comfyui/comfyui.py new file mode 100644 index 0000000000000000000000000000000000000000..a8127dd23f155358813b47848deb102ec15e551d --- /dev/null +++ b/api/core/tools/provider/builtin/comfyui/comfyui.py @@ -0,0 +1,24 @@ +from typing import Any + +import websocket +from yarl import URL + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class ComfyUIProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + ws = websocket.WebSocket() + base_url = URL(credentials.get("base_url")) + ws_protocol = "ws" + if base_url.scheme == "https": + ws_protocol = "wss" + ws_address = f"{ws_protocol}://{base_url.authority}/ws?clientId=test123" + + try: + ws.connect(ws_address) + except Exception as e: + raise ToolProviderCredentialValidationError(f"can not connect to {ws_address}") + finally: + ws.close() diff --git a/api/core/tools/provider/builtin/comfyui/comfyui.yaml b/api/core/tools/provider/builtin/comfyui/comfyui.yaml new file mode 100644 index 0000000000000000000000000000000000000000..24ae43cd44051ece01b51e5739968c9162965fb4 --- /dev/null +++ b/api/core/tools/provider/builtin/comfyui/comfyui.yaml @@ -0,0 +1,23 @@ +identity: + author: Qun + name: comfyui + label: + en_US: ComfyUI + zh_Hans: ComfyUI + description: + en_US: ComfyUI is a tool for generating images which can be deployed locally. + zh_Hans: ComfyUI 是一个可以在本地部署的图片生成的工具。 + icon: icon.png + tags: + - image +credentials_for_provider: + base_url: + type: text-input + required: true + label: + en_US: The URL of ComfyUI Server + zh_Hans: ComfyUI服务器的URL + placeholder: + en_US: Please input your ComfyUI server's Base URL + zh_Hans: 请输入你的 ComfyUI 服务器的 Base URL + url: https://docs.dify.ai/guides/tools/tool-configuration/comfyui diff --git a/api/core/tools/provider/builtin/comfyui/tools/comfyui_client.py b/api/core/tools/provider/builtin/comfyui/tools/comfyui_client.py new file mode 100644 index 0000000000000000000000000000000000000000..2bf10ce8ff263211a7c237f2f0202e066b93e27b --- /dev/null +++ b/api/core/tools/provider/builtin/comfyui/tools/comfyui_client.py @@ -0,0 +1,131 @@ +import json +import random +import uuid + +import httpx +from websocket import WebSocket +from yarl import URL + +from core.file.file_manager import download +from core.file.models import File + + +class ComfyUiClient: + def __init__(self, base_url: str): + self.base_url = URL(base_url) + + def get_history(self, prompt_id: str) -> dict: + res = httpx.get(str(self.base_url / "history"), params={"prompt_id": prompt_id}) + history = res.json()[prompt_id] + return history + + def get_image(self, filename: str, subfolder: str, folder_type: str) -> bytes: + response = httpx.get( + str(self.base_url / "view"), + params={"filename": filename, "subfolder": subfolder, "type": folder_type}, + ) + return response.content + + def upload_image(self, image_file: File) -> dict: + file = download(image_file) + files = {"image": (image_file.filename, file, image_file.mime_type), "overwrite": "true"} + res = httpx.post(str(self.base_url / "upload/image"), files=files) + return res.json() + + def queue_prompt(self, client_id: str, prompt: dict) -> str: + res = httpx.post(str(self.base_url / "prompt"), json={"client_id": client_id, "prompt": prompt}) + prompt_id = res.json()["prompt_id"] + return prompt_id + + def open_websocket_connection(self) -> tuple[WebSocket, str]: + client_id = str(uuid.uuid4()) + ws = WebSocket() + ws_protocol = "ws" + if self.base_url.scheme == "https": + ws_protocol = "wss" + ws_address = f"{ws_protocol}://{self.base_url.authority}/ws?clientId={client_id}" + ws.connect(ws_address) + return ws, client_id + + def set_prompt_by_ksampler(self, origin_prompt: dict, positive_prompt: str, negative_prompt: str = "") -> dict: + prompt = origin_prompt.copy() + id_to_class_type = {id: details["class_type"] for id, details in prompt.items()} + k_sampler = [key for key, value in id_to_class_type.items() if value == "KSampler"][0] + positive_input_id = prompt.get(k_sampler)["inputs"]["positive"][0] + prompt.get(positive_input_id)["inputs"]["text"] = positive_prompt + + if negative_prompt != "": + negative_input_id = prompt.get(k_sampler)["inputs"]["negative"][0] + prompt.get(negative_input_id)["inputs"]["text"] = negative_prompt + + return prompt + + def set_prompt_images_by_ids(self, origin_prompt: dict, image_names: list[str], image_ids: list[str]) -> dict: + prompt = origin_prompt.copy() + for index, image_node_id in enumerate(image_ids): + prompt[image_node_id]["inputs"]["image"] = image_names[index] + return prompt + + def set_prompt_images_by_default(self, origin_prompt: dict, image_names: list[str]) -> dict: + prompt = origin_prompt.copy() + id_to_class_type = {id: details["class_type"] for id, details in prompt.items()} + load_image_nodes = [key for key, value in id_to_class_type.items() if value == "LoadImage"] + for load_image, image_name in zip(load_image_nodes, image_names): + prompt.get(load_image)["inputs"]["image"] = image_name + return prompt + + def set_prompt_seed_by_id(self, origin_prompt: dict, seed_id: str) -> dict: + prompt = origin_prompt.copy() + if seed_id not in prompt: + raise Exception("Not a valid seed node") + if "seed" in prompt[seed_id]["inputs"]: + prompt[seed_id]["inputs"]["seed"] = random.randint(10**14, 10**15 - 1) + elif "noise_seed" in prompt[seed_id]["inputs"]: + prompt[seed_id]["inputs"]["noise_seed"] = random.randint(10**14, 10**15 - 1) + else: + raise Exception("Not a valid seed node") + return prompt + + def track_progress(self, prompt: dict, ws: WebSocket, prompt_id: str): + node_ids = list(prompt.keys()) + finished_nodes = [] + + while True: + out = ws.recv() + if isinstance(out, str): + message = json.loads(out) + if message["type"] == "progress": + data = message["data"] + current_step = data["value"] + print("In K-Sampler -> Step: ", current_step, " of: ", data["max"]) + if message["type"] == "execution_cached": + data = message["data"] + for itm in data["nodes"]: + if itm not in finished_nodes: + finished_nodes.append(itm) + print("Progress: ", len(finished_nodes), "/", len(node_ids), " Tasks done") + if message["type"] == "executing": + data = message["data"] + if data["node"] not in finished_nodes: + finished_nodes.append(data["node"]) + print("Progress: ", len(finished_nodes), "/", len(node_ids), " Tasks done") + + if data["node"] is None and data["prompt_id"] == prompt_id: + break # Execution is done + else: + continue + + def generate_image_by_prompt(self, prompt: dict) -> list[bytes]: + try: + ws, client_id = self.open_websocket_connection() + prompt_id = self.queue_prompt(client_id, prompt) + self.track_progress(prompt, ws, prompt_id) + history = self.get_history(prompt_id) + images = [] + for output in history["outputs"].values(): + for img in output.get("images", []): + image_data = self.get_image(img["filename"], img["subfolder"], img["type"]) + images.append((image_data, img["filename"])) + return images + finally: + ws.close() diff --git a/api/core/tools/provider/builtin/comfyui/tools/comfyui_stable_diffusion.py b/api/core/tools/provider/builtin/comfyui/tools/comfyui_stable_diffusion.py new file mode 100644 index 0000000000000000000000000000000000000000..eaa4b0d02755688a71f378e986a1126a95f149e8 --- /dev/null +++ b/api/core/tools/provider/builtin/comfyui/tools/comfyui_stable_diffusion.py @@ -0,0 +1,475 @@ +import json +import os +import random +import uuid +from copy import deepcopy +from enum import Enum +from typing import Any, Union + +import websocket +from httpx import get, post +from yarl import URL + +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter, ToolParameterOption +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.tool.builtin_tool import BuiltinTool + +SD_TXT2IMG_OPTIONS = {} +LORA_NODE = { + "inputs": {"lora_name": "", "strength_model": 1, "strength_clip": 1, "model": ["11", 0], "clip": ["11", 1]}, + "class_type": "LoraLoader", + "_meta": {"title": "Load LoRA"}, +} +FluxGuidanceNode = { + "inputs": {"guidance": 3.5, "conditioning": ["6", 0]}, + "class_type": "FluxGuidance", + "_meta": {"title": "FluxGuidance"}, +} + + +class ModelType(Enum): + SD15 = 1 + SDXL = 2 + SD3 = 3 + FLUX = 4 + + +class ComfyuiStableDiffusionTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + # base url + base_url = self.runtime.credentials.get("base_url", "") + if not base_url: + return self.create_text_message("Please input base_url") + + if tool_parameters.get("model"): + self.runtime.credentials["model"] = tool_parameters["model"] + + model = self.runtime.credentials.get("model", None) + if not model: + return self.create_text_message("Please input model") + + # prompt + prompt = tool_parameters.get("prompt", "") + if not prompt: + return self.create_text_message("Please input prompt") + + # get negative prompt + negative_prompt = tool_parameters.get("negative_prompt", "") + + # get size + width = tool_parameters.get("width", 1024) + height = tool_parameters.get("height", 1024) + + # get steps + steps = tool_parameters.get("steps", 1) + + # get sampler_name + sampler_name = tool_parameters.get("sampler_name", "euler") + + # scheduler + scheduler = tool_parameters.get("scheduler", "normal") + + # get cfg + cfg = tool_parameters.get("cfg", 7.0) + + # get model type + model_type = tool_parameters.get("model_type", ModelType.SD15.name) + + # get lora + # supports up to 3 loras + lora_list = [] + lora_strength_list = [] + if tool_parameters.get("lora_1"): + lora_list.append(tool_parameters["lora_1"]) + lora_strength_list.append(tool_parameters.get("lora_strength_1", 1)) + if tool_parameters.get("lora_2"): + lora_list.append(tool_parameters["lora_2"]) + lora_strength_list.append(tool_parameters.get("lora_strength_2", 1)) + if tool_parameters.get("lora_3"): + lora_list.append(tool_parameters["lora_3"]) + lora_strength_list.append(tool_parameters.get("lora_strength_3", 1)) + + return self.text2img( + base_url=base_url, + model=model, + model_type=model_type, + prompt=prompt, + negative_prompt=negative_prompt, + width=width, + height=height, + steps=steps, + sampler_name=sampler_name, + scheduler=scheduler, + cfg=cfg, + lora_list=lora_list, + lora_strength_list=lora_strength_list, + ) + + def get_checkpoints(self) -> list[str]: + """ + get checkpoints + """ + try: + base_url = self.runtime.credentials.get("base_url", None) + if not base_url: + return [] + api_url = str(URL(base_url) / "models" / "checkpoints") + response = get(url=api_url, timeout=(2, 10)) + if response.status_code != 200: + return [] + else: + return response.json() + except Exception as e: + return [] + + def get_loras(self) -> list[str]: + """ + get loras + """ + try: + base_url = self.runtime.credentials.get("base_url", None) + if not base_url: + return [] + api_url = str(URL(base_url) / "models" / "loras") + response = get(url=api_url, timeout=(2, 10)) + if response.status_code != 200: + return [] + else: + return response.json() + except Exception as e: + return [] + + def get_sample_methods(self) -> tuple[list[str], list[str]]: + """ + get sample method + """ + try: + base_url = self.runtime.credentials.get("base_url", None) + if not base_url: + return [], [] + api_url = str(URL(base_url) / "object_info" / "KSampler") + response = get(url=api_url, timeout=(2, 10)) + if response.status_code != 200: + return [], [] + else: + data = response.json()["KSampler"]["input"]["required"] + return data["sampler_name"][0], data["scheduler"][0] + except Exception as e: + return [], [] + + def validate_models(self) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + validate models + """ + try: + base_url = self.runtime.credentials.get("base_url", None) + if not base_url: + raise ToolProviderCredentialValidationError("Please input base_url") + model = self.runtime.credentials.get("model", None) + if not model: + raise ToolProviderCredentialValidationError("Please input model") + + api_url = str(URL(base_url) / "models" / "checkpoints") + response = get(url=api_url, timeout=(2, 10)) + if response.status_code != 200: + raise ToolProviderCredentialValidationError("Failed to get models") + else: + models = response.json() + if len([d for d in models if d == model]) > 0: + return self.create_text_message(json.dumps(models)) + else: + raise ToolProviderCredentialValidationError(f"model {model} does not exist") + except Exception as e: + raise ToolProviderCredentialValidationError(f"Failed to get models, {e}") + + def get_history(self, base_url, prompt_id): + """ + get history + """ + url = str(URL(base_url) / "history") + respond = get(url, params={"prompt_id": prompt_id}, timeout=(2, 10)) + return respond.json() + + def download_image(self, base_url, filename, subfolder, folder_type): + """ + download image + """ + url = str(URL(base_url) / "view") + response = get(url, params={"filename": filename, "subfolder": subfolder, "type": folder_type}, timeout=(2, 10)) + return response.content + + def queue_prompt_image(self, base_url, client_id, prompt): + """ + send prompt task and rotate + """ + # initiate task execution + url = str(URL(base_url) / "prompt") + respond = post(url, data=json.dumps({"client_id": client_id, "prompt": prompt}), timeout=(2, 10)) + prompt_id = respond.json()["prompt_id"] + + ws = websocket.WebSocket() + if "https" in base_url: + ws_url = base_url.replace("https", "ws") + else: + ws_url = base_url.replace("http", "ws") + ws.connect(str(URL(f"{ws_url}") / "ws") + f"?clientId={client_id}", timeout=120) + + # websocket rotate execution status + output_images = {} + while True: + out = ws.recv() + if isinstance(out, str): + message = json.loads(out) + if message["type"] == "executing": + data = message["data"] + if data["node"] is None and data["prompt_id"] == prompt_id: + break # Execution is done + elif message["type"] == "status": + data = message["data"] + if data["status"]["exec_info"]["queue_remaining"] == 0 and data.get("sid"): + break # Execution is done + else: + continue # previews are binary data + + # download image when execution finished + history = self.get_history(base_url, prompt_id)[prompt_id] + for o in history["outputs"]: + for node_id in history["outputs"]: + node_output = history["outputs"][node_id] + if "images" in node_output: + images_output = [] + for image in node_output["images"]: + image_data = self.download_image(base_url, image["filename"], image["subfolder"], image["type"]) + images_output.append(image_data) + output_images[node_id] = images_output + + ws.close() + + return output_images + + def text2img( + self, + base_url: str, + model: str, + model_type: str, + prompt: str, + negative_prompt: str, + width: int, + height: int, + steps: int, + sampler_name: str, + scheduler: str, + cfg: float, + lora_list: list, + lora_strength_list: list, + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + generate image + """ + if not SD_TXT2IMG_OPTIONS: + current_dir = os.path.dirname(os.path.realpath(__file__)) + with open(os.path.join(current_dir, "txt2img.json")) as file: + SD_TXT2IMG_OPTIONS.update(json.load(file)) + + draw_options = deepcopy(SD_TXT2IMG_OPTIONS) + draw_options["3"]["inputs"]["steps"] = steps + draw_options["3"]["inputs"]["sampler_name"] = sampler_name + draw_options["3"]["inputs"]["scheduler"] = scheduler + draw_options["3"]["inputs"]["cfg"] = cfg + # generate different image when using same prompt next time + draw_options["3"]["inputs"]["seed"] = random.randint(0, 100000000) + draw_options["4"]["inputs"]["ckpt_name"] = model + draw_options["5"]["inputs"]["width"] = width + draw_options["5"]["inputs"]["height"] = height + draw_options["6"]["inputs"]["text"] = prompt + draw_options["7"]["inputs"]["text"] = negative_prompt + # if the model is SD3 or FLUX series, the Latent class should be corresponding to SD3 Latent + if model_type in {ModelType.SD3.name, ModelType.FLUX.name}: + draw_options["5"]["class_type"] = "EmptySD3LatentImage" + + if lora_list: + # last Lora node link to KSampler node + draw_options["3"]["inputs"]["model"][0] = "10" + # last Lora node link to positive and negative Clip node + draw_options["6"]["inputs"]["clip"][0] = "10" + draw_options["7"]["inputs"]["clip"][0] = "10" + # every Lora node link to next Lora node, and Checkpoints node link to first Lora node + for i, (lora, strength) in enumerate(zip(lora_list, lora_strength_list), 10): + if i - 10 == len(lora_list) - 1: + next_node_id = "4" + else: + next_node_id = str(i + 1) + lora_node = deepcopy(LORA_NODE) + lora_node["inputs"]["lora_name"] = lora + lora_node["inputs"]["strength_model"] = strength + lora_node["inputs"]["strength_clip"] = strength + lora_node["inputs"]["model"][0] = next_node_id + lora_node["inputs"]["clip"][0] = next_node_id + draw_options[str(i)] = lora_node + + # FLUX need to add FluxGuidance Node + if model_type == ModelType.FLUX.name: + last_node_id = str(10 + len(lora_list)) + draw_options[last_node_id] = deepcopy(FluxGuidanceNode) + draw_options[last_node_id]["inputs"]["conditioning"][0] = "6" + draw_options["3"]["inputs"]["positive"][0] = last_node_id + + try: + client_id = str(uuid.uuid4()) + result = self.queue_prompt_image(base_url, client_id, prompt=draw_options) + + # get first image + image = b"" + for node in result: + for img in result[node]: + if img: + image = img + break + + return self.create_blob_message( + blob=image, meta={"mime_type": "image/png"}, save_as=self.VariableKey.IMAGE.value + ) + + except Exception as e: + return self.create_text_message(f"Failed to generate image: {str(e)}") + + def get_runtime_parameters(self) -> list[ToolParameter]: + parameters = [ + ToolParameter( + name="prompt", + label=I18nObject(en_US="Prompt", zh_Hans="Prompt"), + human_description=I18nObject( + en_US="Image prompt, you can check the official documentation of Stable Diffusion", + zh_Hans="图像提示词,您可以查看 Stable Diffusion 的官方文档", + ), + type=ToolParameter.ToolParameterType.STRING, + form=ToolParameter.ToolParameterForm.LLM, + llm_description="Image prompt of Stable Diffusion, you should describe the image " + "you want to generate as a list of words as possible as detailed, " + "the prompt must be written in English.", + required=True, + ), + ] + if self.runtime.credentials: + try: + models = self.get_checkpoints() + if len(models) != 0: + parameters.append( + ToolParameter( + name="model", + label=I18nObject(en_US="Model", zh_Hans="Model"), + human_description=I18nObject( + en_US="Model of Stable Diffusion or FLUX, " + "you can check the official documentation of Stable Diffusion or FLUX", + zh_Hans="Stable Diffusion 或者 FLUX 的模型,您可以查看 Stable Diffusion 的官方文档", + ), + type=ToolParameter.ToolParameterType.SELECT, + form=ToolParameter.ToolParameterForm.FORM, + llm_description="Model of Stable Diffusion or FLUX, " + "you can check the official documentation of Stable Diffusion or FLUX", + required=True, + default=models[0], + options=[ + ToolParameterOption(value=i, label=I18nObject(en_US=i, zh_Hans=i)) for i in models + ], + ) + ) + loras = self.get_loras() + if len(loras) != 0: + for n in range(1, 4): + parameters.append( + ToolParameter( + name=f"lora_{n}", + label=I18nObject(en_US=f"Lora {n}", zh_Hans=f"Lora {n}"), + human_description=I18nObject( + en_US="Lora of Stable Diffusion, " + "you can check the official documentation of Stable Diffusion", + zh_Hans="Stable Diffusion 的 Lora 模型,您可以查看 Stable Diffusion 的官方文档", + ), + type=ToolParameter.ToolParameterType.SELECT, + form=ToolParameter.ToolParameterForm.FORM, + llm_description="Lora of Stable Diffusion, " + "you can check the official documentation of " + "Stable Diffusion", + required=False, + options=[ + ToolParameterOption(value=i, label=I18nObject(en_US=i, zh_Hans=i)) for i in loras + ], + ) + ) + sample_methods, schedulers = self.get_sample_methods() + if len(sample_methods) != 0: + parameters.append( + ToolParameter( + name="sampler_name", + label=I18nObject(en_US="Sampling method", zh_Hans="Sampling method"), + human_description=I18nObject( + en_US="Sampling method of Stable Diffusion, " + "you can check the official documentation of Stable Diffusion", + zh_Hans="Stable Diffusion 的Sampling method,您可以查看 Stable Diffusion 的官方文档", + ), + type=ToolParameter.ToolParameterType.SELECT, + form=ToolParameter.ToolParameterForm.FORM, + llm_description="Sampling method of Stable Diffusion, " + "you can check the official documentation of Stable Diffusion", + required=True, + default=sample_methods[0], + options=[ + ToolParameterOption(value=i, label=I18nObject(en_US=i, zh_Hans=i)) + for i in sample_methods + ], + ) + ) + if len(schedulers) != 0: + parameters.append( + ToolParameter( + name="scheduler", + label=I18nObject(en_US="Scheduler", zh_Hans="Scheduler"), + human_description=I18nObject( + en_US="Scheduler of Stable Diffusion, " + "you can check the official documentation of Stable Diffusion", + zh_Hans="Stable Diffusion 的Scheduler,您可以查看 Stable Diffusion 的官方文档", + ), + type=ToolParameter.ToolParameterType.SELECT, + form=ToolParameter.ToolParameterForm.FORM, + llm_description="Scheduler of Stable Diffusion, " + "you can check the official documentation of Stable Diffusion", + required=True, + default=schedulers[0], + options=[ + ToolParameterOption(value=i, label=I18nObject(en_US=i, zh_Hans=i)) for i in schedulers + ], + ) + ) + parameters.append( + ToolParameter( + name="model_type", + label=I18nObject(en_US="Model Type", zh_Hans="Model Type"), + human_description=I18nObject( + en_US="Model Type of Stable Diffusion or Flux, " + "you can check the official documentation of Stable Diffusion or Flux", + zh_Hans="Stable Diffusion 或 FLUX 的模型类型," + "您可以查看 Stable Diffusion 或 Flux 的官方文档", + ), + type=ToolParameter.ToolParameterType.SELECT, + form=ToolParameter.ToolParameterForm.FORM, + llm_description="Model Type of Stable Diffusion or Flux, " + "you can check the official documentation of Stable Diffusion or Flux", + required=True, + default=ModelType.SD15.name, + options=[ + ToolParameterOption(value=i, label=I18nObject(en_US=i, zh_Hans=i)) + for i in ModelType.__members__ + ], + ) + ) + except: + pass + + return parameters diff --git a/api/core/tools/provider/builtin/comfyui/tools/comfyui_stable_diffusion.yaml b/api/core/tools/provider/builtin/comfyui/tools/comfyui_stable_diffusion.yaml new file mode 100644 index 0000000000000000000000000000000000000000..75fe746965196a650bfe030153608c37067c1222 --- /dev/null +++ b/api/core/tools/provider/builtin/comfyui/tools/comfyui_stable_diffusion.yaml @@ -0,0 +1,212 @@ +identity: + name: txt2img + author: Qun + label: + en_US: Txt2Img + zh_Hans: Txt2Img + pt_BR: Txt2Img +description: + human: + en_US: a pre-defined comfyui workflow that can use one model and up to 3 loras to generate images. Support SD1.5, SDXL, SD3 and FLUX which contain text encoders/clip, but does not support models that requires a triple clip loader. + zh_Hans: 一个预定义的 ComfyUI 工作流,可以使用一个模型和最多3个loras来生成图像。支持包含文本编码器/clip的SD1.5、SDXL、SD3和FLUX,但不支持需要clip加载器的模型。 + pt_BR: a pre-defined comfyui workflow that can use one model and up to 3 loras to generate images. Support SD1.5, SDXL, SD3 and FLUX which contain text encoders/clip, but does not support models that requires a triple clip loader. + llm: draw the image you want based on your prompt. +parameters: + - name: prompt + type: string + required: true + label: + en_US: Prompt + zh_Hans: 提示词 + pt_BR: Prompt + human_description: + en_US: Image prompt, you can check the official documentation of Stable Diffusion or FLUX + zh_Hans: 图像提示词,您可以查看 Stable Diffusion 或者 FLUX 的官方文档 + pt_BR: Image prompt, you can check the official documentation of Stable Diffusion or FLUX + llm_description: Image prompt of Stable Diffusion, you should describe the image you want to generate as a list of words as possible as detailed, the prompt must be written in English. + form: llm + - name: model + type: string + required: true + label: + en_US: Model Name + zh_Hans: 模型名称 + pt_BR: Model Name + human_description: + en_US: Model Name + zh_Hans: 模型名称 + pt_BR: Model Name + form: form + - name: model_type + type: string + required: true + label: + en_US: Model Type + zh_Hans: 模型类型 + pt_BR: Model Type + human_description: + en_US: Model Type + zh_Hans: 模型类型 + pt_BR: Model Type + form: form + - name: lora_1 + type: string + required: false + label: + en_US: Lora 1 + zh_Hans: Lora 1 + pt_BR: Lora 1 + human_description: + en_US: Lora 1 + zh_Hans: Lora 1 + pt_BR: Lora 1 + form: form + - name: lora_strength_1 + type: number + required: false + label: + en_US: Lora Strength 1 + zh_Hans: Lora Strength 1 + pt_BR: Lora Strength 1 + human_description: + en_US: Lora Strength 1 + zh_Hans: Lora模型的权重 + pt_BR: Lora Strength 1 + form: form + - name: steps + type: number + required: false + label: + en_US: Steps + zh_Hans: Steps + pt_BR: Steps + human_description: + en_US: Steps + zh_Hans: Steps + pt_BR: Steps + form: form + default: 20 + - name: width + type: number + required: false + label: + en_US: Width + zh_Hans: Width + pt_BR: Width + human_description: + en_US: Width + zh_Hans: Width + pt_BR: Width + form: form + default: 1024 + - name: height + type: number + required: false + label: + en_US: Height + zh_Hans: Height + pt_BR: Height + human_description: + en_US: Height + zh_Hans: Height + pt_BR: Height + form: form + default: 1024 + - name: negative_prompt + type: string + required: false + label: + en_US: Negative prompt + zh_Hans: Negative prompt + pt_BR: Negative prompt + human_description: + en_US: Negative prompt + zh_Hans: Negative prompt + pt_BR: Negative prompt + form: form + default: bad art, ugly, deformed, watermark, duplicated, discontinuous lines + - name: cfg + type: number + required: false + label: + en_US: CFG Scale + zh_Hans: CFG Scale + pt_BR: CFG Scale + human_description: + en_US: CFG Scale + zh_Hans: 提示词相关性(CFG Scale) + pt_BR: CFG Scale + form: form + default: 7.0 + - name: sampler_name + type: string + required: false + label: + en_US: Sampling method + zh_Hans: Sampling method + pt_BR: Sampling method + human_description: + en_US: Sampling method + zh_Hans: Sampling method + pt_BR: Sampling method + form: form + - name: scheduler + type: string + required: false + label: + en_US: Scheduler + zh_Hans: Scheduler + pt_BR: Scheduler + human_description: + en_US: Scheduler + zh_Hans: Scheduler + pt_BR: Scheduler + form: form + - name: lora_2 + type: string + required: false + label: + en_US: Lora 2 + zh_Hans: Lora 2 + pt_BR: Lora 2 + human_description: + en_US: Lora 2 + zh_Hans: Lora 2 + pt_BR: Lora 2 + form: form + - name: lora_strength_2 + type: number + required: false + label: + en_US: Lora Strength 2 + zh_Hans: Lora Strength 2 + pt_BR: Lora Strength 2 + human_description: + en_US: Lora Strength 2 + zh_Hans: Lora模型的权重 + pt_BR: Lora Strength 2 + form: form + - name: lora_3 + type: string + required: false + label: + en_US: Lora 3 + zh_Hans: Lora 3 + pt_BR: Lora 3 + human_description: + en_US: Lora 3 + zh_Hans: Lora 3 + pt_BR: Lora 3 + form: form + - name: lora_strength_3 + type: number + required: false + label: + en_US: Lora Strength 3 + zh_Hans: Lora Strength 3 + pt_BR: Lora Strength 3 + human_description: + en_US: Lora Strength 3 + zh_Hans: Lora模型的权重 + pt_BR: Lora Strength 3 + form: form diff --git a/api/core/tools/provider/builtin/comfyui/tools/comfyui_workflow.py b/api/core/tools/provider/builtin/comfyui/tools/comfyui_workflow.py new file mode 100644 index 0000000000000000000000000000000000000000..eb085f221ebdda6f48c72bdfb292bde29652de95 --- /dev/null +++ b/api/core/tools/provider/builtin/comfyui/tools/comfyui_workflow.py @@ -0,0 +1,87 @@ +import json +import mimetypes +from typing import Any + +from core.file import FileType +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.errors import ToolParameterValidationError +from core.tools.provider.builtin.comfyui.tools.comfyui_client import ComfyUiClient +from core.tools.tool.builtin_tool import BuiltinTool + + +def sanitize_json_string(s): + escape_dict = { + "\n": "\\n", + "\r": "\\r", + "\t": "\\t", + "\b": "\\b", + "\f": "\\f", + } + for char, escaped in escape_dict.items(): + s = s.replace(char, escaped) + + return s + + +class ComfyUIWorkflowTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage | list[ToolInvokeMessage]: + comfyui = ComfyUiClient(self.runtime.credentials["base_url"]) + + positive_prompt = tool_parameters.get("positive_prompt", "") + negative_prompt = tool_parameters.get("negative_prompt", "") + images = tool_parameters.get("images") or [] + workflow = tool_parameters.get("workflow_json") + image_names = [] + for image in images: + if image.type != FileType.IMAGE: + continue + image_name = comfyui.upload_image(image).get("name") + image_names.append(image_name) + + set_prompt_with_ksampler = True + if "{{positive_prompt}}" in workflow: + set_prompt_with_ksampler = False + workflow = workflow.replace("{{positive_prompt}}", positive_prompt.replace('"', "'")) + workflow = workflow.replace("{{negative_prompt}}", negative_prompt.replace('"', "'")) + + try: + prompt = json.loads(workflow) + except json.JSONDecodeError: + cleaned_string = sanitize_json_string(workflow) + try: + prompt = json.loads(cleaned_string) + except: + return self.create_text_message("the Workflow JSON is not correct") + + if set_prompt_with_ksampler: + try: + prompt = comfyui.set_prompt_by_ksampler(prompt, positive_prompt, negative_prompt) + except: + raise ToolParameterValidationError( + "Failed set prompt with KSampler, try replace prompt to {{positive_prompt}} in your workflow json" + ) + + if image_names: + if image_ids := tool_parameters.get("image_ids"): + image_ids = image_ids.split(",") + try: + prompt = comfyui.set_prompt_images_by_ids(prompt, image_names, image_ids) + except: + raise ToolParameterValidationError("the Image Node ID List not match your upload image files.") + else: + prompt = comfyui.set_prompt_images_by_default(prompt, image_names) + + if seed_id := tool_parameters.get("seed_id"): + prompt = comfyui.set_prompt_seed_by_id(prompt, seed_id) + + images = comfyui.generate_image_by_prompt(prompt) + result = [] + for image_data, filename in images: + result.append( + self.create_blob_message( + blob=image_data, + meta={"mime_type": mimetypes.guess_type(filename)[0]}, + save_as=self.VariableKey.IMAGE.value, + ) + ) + return result diff --git a/api/core/tools/provider/builtin/comfyui/tools/comfyui_workflow.yaml b/api/core/tools/provider/builtin/comfyui/tools/comfyui_workflow.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9428acbe9436423591ceb1604ea5914039aa64e7 --- /dev/null +++ b/api/core/tools/provider/builtin/comfyui/tools/comfyui_workflow.yaml @@ -0,0 +1,63 @@ +identity: + name: workflow + author: hjlarry + label: + en_US: workflow + zh_Hans: 工作流 +description: + human: + en_US: Run ComfyUI workflow. + zh_Hans: 运行ComfyUI工作流。 + llm: Run ComfyUI workflow. +parameters: + - name: positive_prompt + type: string + label: + en_US: Prompt + zh_Hans: 提示词 + llm_description: Image prompt, you should describe the image you want to generate as a list of words as possible as detailed, the prompt must be written in English. + form: llm + - name: negative_prompt + type: string + label: + en_US: Negative Prompt + zh_Hans: 负面提示词 + llm_description: Negative prompt, you should describe the image you don't want to generate as a list of words as possible as detailed, the prompt must be written in English. + form: llm + - name: images + type: files + label: + en_US: Input Images + zh_Hans: 输入的图片 + llm_description: The input images, used to transfer to the comfyui workflow to generate another image. + form: llm + - name: workflow_json + type: string + required: true + label: + en_US: Workflow JSON + human_description: + en_US: exported from ComfyUI workflow + zh_Hans: 从ComfyUI的工作流中导出 + form: form + - name: image_ids + type: string + label: + en_US: Image Node ID List + zh_Hans: 图片节点ID列表 + placeholder: + en_US: Use commas to separate multiple node ID + zh_Hans: 多个节点ID时使用半角逗号分隔 + human_description: + en_US: When the workflow has multiple image nodes, enter the ID list of these nodes, and the images will be passed to ComfyUI in the order of the list. + zh_Hans: 当工作流有多个图片节点时,输入这些节点的ID列表,图片将按列表顺序传给ComfyUI + form: form + - name: seed_id + type: string + label: + en_US: Seed Node Id + zh_Hans: 种子节点ID + human_description: + en_US: If you need to generate different images each time, you need to enter the ID of the seed node. + zh_Hans: 如果需要每次生成时使用不同的种子,需要输入包含种子的节点的ID + form: form diff --git a/api/core/tools/provider/builtin/comfyui/tools/txt2img.json b/api/core/tools/provider/builtin/comfyui/tools/txt2img.json new file mode 100644 index 0000000000000000000000000000000000000000..8ea869ff106c3827c3fd79fca1d9090ff17cd6a9 --- /dev/null +++ b/api/core/tools/provider/builtin/comfyui/tools/txt2img.json @@ -0,0 +1,107 @@ +{ + "3": { + "inputs": { + "seed": 156680208700286, + "steps": 20, + "cfg": 8, + "sampler_name": "euler", + "scheduler": "normal", + "denoise": 1, + "model": [ + "4", + 0 + ], + "positive": [ + "6", + 0 + ], + "negative": [ + "7", + 0 + ], + "latent_image": [ + "5", + 0 + ] + }, + "class_type": "KSampler", + "_meta": { + "title": "KSampler" + } + }, + "4": { + "inputs": { + "ckpt_name": "3dAnimationDiffusion_v10.safetensors" + }, + "class_type": "CheckpointLoaderSimple", + "_meta": { + "title": "Load Checkpoint" + } + }, + "5": { + "inputs": { + "width": 512, + "height": 512, + "batch_size": 1 + }, + "class_type": "EmptyLatentImage", + "_meta": { + "title": "Empty Latent Image" + } + }, + "6": { + "inputs": { + "text": "beautiful scenery nature glass bottle landscape, , purple galaxy bottle,", + "clip": [ + "4", + 1 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "CLIP Text Encode (Prompt)" + } + }, + "7": { + "inputs": { + "text": "text, watermark", + "clip": [ + "4", + 1 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "CLIP Text Encode (Prompt)" + } + }, + "8": { + "inputs": { + "samples": [ + "3", + 0 + ], + "vae": [ + "4", + 2 + ] + }, + "class_type": "VAEDecode", + "_meta": { + "title": "VAE Decode" + } + }, + "9": { + "inputs": { + "filename_prefix": "ComfyUI", + "images": [ + "8", + 0 + ] + }, + "class_type": "SaveImage", + "_meta": { + "title": "Save Image" + } + } +} \ No newline at end of file diff --git a/api/core/tools/provider/builtin/crossref/_assets/icon.svg b/api/core/tools/provider/builtin/crossref/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..aa629de7cb16605423302c6dd4f9bcdf11a27ef3 --- /dev/null +++ b/api/core/tools/provider/builtin/crossref/_assets/icon.svg @@ -0,0 +1,49 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/api/core/tools/provider/builtin/crossref/crossref.py b/api/core/tools/provider/builtin/crossref/crossref.py new file mode 100644 index 0000000000000000000000000000000000000000..8ba3c1b48ae6d7a8f86b96e0cfb54abd2e827e2e --- /dev/null +++ b/api/core/tools/provider/builtin/crossref/crossref.py @@ -0,0 +1,20 @@ +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.crossref.tools.query_doi import CrossRefQueryDOITool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class CrossRefProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + try: + CrossRefQueryDOITool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "doi": "10.1007/s00894-022-05373-8", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/crossref/crossref.yaml b/api/core/tools/provider/builtin/crossref/crossref.yaml new file mode 100644 index 0000000000000000000000000000000000000000..da67fbec3a480ba1f2a58aff70b217f741f62118 --- /dev/null +++ b/api/core/tools/provider/builtin/crossref/crossref.yaml @@ -0,0 +1,29 @@ +identity: + author: Sakura4036 + name: crossref + label: + en_US: CrossRef + zh_Hans: CrossRef + description: + en_US: Crossref is a cross-publisher reference linking registration query system using DOI technology created in 2000. Crossref establishes cross-database links between the reference list and citation full text of papers, making it very convenient for readers to access the full text of papers. + zh_Hans: Crossref是于2000年创建的使用DOI技术的跨出版商参考文献链接注册查询系统。Crossref建立了在论文的参考文献列表和引文全文之间的跨数据库链接,使得读者能够非常便捷地获取文献全文。 + icon: icon.svg + tags: + - search +credentials_for_provider: + mailto: + type: text-input + required: true + label: + en_US: email address + zh_Hans: email地址 + pt_BR: email address + placeholder: + en_US: Please input your email address + zh_Hans: 请输入你的email地址 + pt_BR: Please input your email address + help: + en_US: According to the requirements of Crossref, an email address is required + zh_Hans: 根据Crossref的要求,需要提供一个邮箱地址 + pt_BR: According to the requirements of Crossref, an email address is required + url: https://api.crossref.org/swagger-ui/index.html diff --git a/api/core/tools/provider/builtin/crossref/tools/query_doi.py b/api/core/tools/provider/builtin/crossref/tools/query_doi.py new file mode 100644 index 0000000000000000000000000000000000000000..746139dd69d27b55d49a0e3daa65f3ac6a70e8d6 --- /dev/null +++ b/api/core/tools/provider/builtin/crossref/tools/query_doi.py @@ -0,0 +1,28 @@ +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.errors import ToolParameterValidationError +from core.tools.tool.builtin_tool import BuiltinTool + + +class CrossRefQueryDOITool(BuiltinTool): + """ + Tool for querying the metadata of a publication using its DOI. + """ + + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + doi = tool_parameters.get("doi") + if not doi: + raise ToolParameterValidationError("doi is required.") + # doc: https://github.com/CrossRef/rest-api-doc + url = f"https://api.crossref.org/works/{doi}" + response = requests.get(url) + response.raise_for_status() + response = response.json() + message = response.get("message", {}) + + return self.create_json_message(message) diff --git a/api/core/tools/provider/builtin/crossref/tools/query_doi.yaml b/api/core/tools/provider/builtin/crossref/tools/query_doi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9c16da25edf2b39aef052bf0c668e940578198fa --- /dev/null +++ b/api/core/tools/provider/builtin/crossref/tools/query_doi.yaml @@ -0,0 +1,23 @@ +identity: + name: crossref_query_doi + author: Sakura4036 + label: + en_US: CrossRef Query DOI + zh_Hans: CrossRef DOI 查询 + pt_BR: CrossRef Query DOI +description: + human: + en_US: A tool for searching literature information using CrossRef by DOI. + zh_Hans: 一个使用CrossRef通过DOI获取文献信息的工具。 + pt_BR: A tool for searching literature information using CrossRef by DOI. + llm: A tool for searching literature information using CrossRef by DOI. +parameters: + - name: doi + type: string + required: true + label: + en_US: DOI + zh_Hans: DOI + pt_BR: DOI + llm_description: DOI for searching in CrossRef + form: llm diff --git a/api/core/tools/provider/builtin/crossref/tools/query_title.py b/api/core/tools/provider/builtin/crossref/tools/query_title.py new file mode 100644 index 0000000000000000000000000000000000000000..e245238183293844c496d9133b2d89936e9615b5 --- /dev/null +++ b/api/core/tools/provider/builtin/crossref/tools/query_title.py @@ -0,0 +1,143 @@ +import time +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +def convert_time_str_to_seconds(time_str: str) -> int: + """ + Convert a time string to seconds. + example: 1s -> 1, 1m30s -> 90, 1h30m -> 5400, 1h30m30s -> 5430 + """ + time_str = time_str.lower().strip().replace(" ", "") + seconds = 0 + if "h" in time_str: + hours, time_str = time_str.split("h") + seconds += int(hours) * 3600 + if "m" in time_str: + minutes, time_str = time_str.split("m") + seconds += int(minutes) * 60 + if "s" in time_str: + seconds += int(time_str.replace("s", "")) + return seconds + + +class CrossRefQueryTitleAPI: + """ + Tool for querying the metadata of a publication using its title. + Crossref API doc: https://github.com/CrossRef/rest-api-doc + """ + + query_url_template: str = "https://api.crossref.org/works?query.bibliographic={query}&rows={rows}&offset={offset}&sort={sort}&order={order}&mailto={mailto}" + rate_limit: int = 50 + rate_interval: float = 1 + max_limit: int = 1000 + + def __init__(self, mailto: str): + self.mailto = mailto + + def _query( + self, + query: str, + rows: int = 5, + offset: int = 0, + sort: str = "relevance", + order: str = "desc", + fuzzy_query: bool = False, + ) -> list[dict]: + """ + Query the metadata of a publication using its title. + :param query: the title of the publication + :param rows: the number of results to return + :param sort: the sort field + :param order: the sort order + :param fuzzy_query: whether to return all items that match the query + """ + url = self.query_url_template.format( + query=query, rows=rows, offset=offset, sort=sort, order=order, mailto=self.mailto + ) + response = requests.get(url) + response.raise_for_status() + rate_limit = int(response.headers["x-ratelimit-limit"]) + # convert time string to seconds + rate_interval = convert_time_str_to_seconds(response.headers["x-ratelimit-interval"]) + + self.rate_limit = rate_limit + self.rate_interval = rate_interval + + response = response.json() + if response["status"] != "ok": + return [] + + message = response["message"] + if fuzzy_query: + # fuzzy query return all items + return message["items"] + else: + for paper in message["items"]: + title = paper["title"][0] + if title.lower() != query.lower(): + continue + return [paper] + return [] + + def query( + self, query: str, rows: int = 5, sort: str = "relevance", order: str = "desc", fuzzy_query: bool = False + ) -> list[dict]: + """ + Query the metadata of a publication using its title. + :param query: the title of the publication + :param rows: the number of results to return + :param sort: the sort field + :param order: the sort order + :param fuzzy_query: whether to return all items that match the query + """ + rows = min(rows, self.max_limit) + if rows > self.rate_limit: + # query multiple times + query_times = rows // self.rate_limit + 1 + results = [] + + for i in range(query_times): + result = self._query( + query, + rows=self.rate_limit, + offset=i * self.rate_limit, + sort=sort, + order=order, + fuzzy_query=fuzzy_query, + ) + if fuzzy_query: + results.extend(result) + else: + # fuzzy_query=False, only one result + if result: + return result + time.sleep(self.rate_interval) + return results + else: + # query once + return self._query(query, rows, sort=sort, order=order, fuzzy_query=fuzzy_query) + + +class CrossRefQueryTitleTool(BuiltinTool): + """ + Tool for querying the metadata of a publication using its title. + """ + + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + query = tool_parameters.get("query") + fuzzy_query = tool_parameters.get("fuzzy_query", False) + rows = tool_parameters.get("rows", 3) + sort = tool_parameters.get("sort", "relevance") + order = tool_parameters.get("order", "desc") + mailto = self.runtime.credentials["mailto"] + + result = CrossRefQueryTitleAPI(mailto).query(query, rows, sort, order, fuzzy_query) + + return [self.create_json_message(r) for r in result] diff --git a/api/core/tools/provider/builtin/crossref/tools/query_title.yaml b/api/core/tools/provider/builtin/crossref/tools/query_title.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5579c77f5293d348e8edca76826f03def8b74887 --- /dev/null +++ b/api/core/tools/provider/builtin/crossref/tools/query_title.yaml @@ -0,0 +1,105 @@ +identity: + name: crossref_query_title + author: Sakura4036 + label: + en_US: CrossRef Title Query + zh_Hans: CrossRef 标题查询 + pt_BR: CrossRef Title Query +description: + human: + en_US: A tool for querying literature information using CrossRef by title. + zh_Hans: 一个使用CrossRef通过标题搜索文献信息的工具。 + pt_BR: A tool for querying literature information using CrossRef by title. + llm: A tool for querying literature information using CrossRef by title. +parameters: + - name: query + type: string + required: true + label: + en_US: 标题 + zh_Hans: 查询语句 + pt_BR: 标题 + human_description: + en_US: Query bibliographic information, useful for citation look up. Includes titles, authors, ISSNs and publication years + zh_Hans: 用于搜索文献信息,有助于查找引用。包括标题,作者,ISSN和出版年份 + pt_BR: Query bibliographic information, useful for citation look up. Includes titles, authors, ISSNs and publication years + llm_description: key words for querying in Web of Science + form: llm + - name: fuzzy_query + type: boolean + default: false + label: + en_US: Whether to fuzzy search + zh_Hans: 是否模糊搜索 + pt_BR: Whether to fuzzy search + human_description: + en_US: used for selecting the query type, fuzzy query returns more results, precise query returns 1 or none + zh_Hans: 用于选择搜索类型,模糊搜索返回更多结果,精确搜索返回1条结果或无 + pt_BR: used for selecting the query type, fuzzy query returns more results, precise query returns 1 or none + form: form + - name: limit + type: number + required: false + label: + en_US: max query number + zh_Hans: 最大搜索数 + pt_BR: max query number + human_description: + en_US: max query number(fuzzy search returns the maximum number of results or precise search the maximum number of matches) + zh_Hans: 最大搜索数(模糊搜索返回的最大结果数或精确搜索最大匹配数) + pt_BR: max query number(fuzzy search returns the maximum number of results or precise search the maximum number of matches) + form: llm + default: 50 + - name: sort + type: select + required: true + options: + - value: relevance + label: + en_US: relevance + zh_Hans: 相关性 + pt_BR: relevance + - value: published + label: + en_US: publication date + zh_Hans: 出版日期 + pt_BR: publication date + - value: references-count + label: + en_US: references-count + zh_Hans: 引用次数 + pt_BR: references-count + default: relevance + label: + en_US: sorting field + zh_Hans: 排序字段 + pt_BR: sorting field + human_description: + en_US: Sorting of query results + zh_Hans: 检索结果的排序字段 + pt_BR: Sorting of query results + form: form + - name: order + type: select + required: true + options: + - value: desc + label: + en_US: descending + zh_Hans: 降序 + pt_BR: descending + - value: asc + label: + en_US: ascending + zh_Hans: 升序 + pt_BR: ascending + default: desc + label: + en_US: Order + zh_Hans: 排序 + pt_BR: Order + human_description: + en_US: Order of query results + zh_Hans: 检索结果的排序方式 + pt_BR: Order of query results + form: form diff --git a/api/core/tools/provider/builtin/dalle/__init__.py b/api/core/tools/provider/builtin/dalle/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/tools/provider/builtin/dalle/_assets/icon.png b/api/core/tools/provider/builtin/dalle/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..b953a7b79e4c4da8228ca9faf763f7bc82366213 --- /dev/null +++ b/api/core/tools/provider/builtin/dalle/_assets/icon.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2b2a3717dc9d6e649b9130ab96722a63d45d54ebac1a306a0bffb3c9edda87b7 +size 156474 diff --git a/api/core/tools/provider/builtin/dalle/dalle.py b/api/core/tools/provider/builtin/dalle/dalle.py new file mode 100644 index 0000000000000000000000000000000000000000..5bd16e49e85e299fdcc4719b1ba56a472f963371 --- /dev/null +++ b/api/core/tools/provider/builtin/dalle/dalle.py @@ -0,0 +1,20 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.dalle.tools.dalle2 import DallE2Tool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class DALLEProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + DallE2Tool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={"prompt": "cute girl, blue eyes, white hair, anime style", "size": "small", "n": 1}, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/dalle/dalle.yaml b/api/core/tools/provider/builtin/dalle/dalle.yaml new file mode 100644 index 0000000000000000000000000000000000000000..37cf93c28aae58bfac96a7f63e3777ce1f491216 --- /dev/null +++ b/api/core/tools/provider/builtin/dalle/dalle.yaml @@ -0,0 +1,61 @@ +identity: + author: Dify + name: dalle + label: + en_US: DALL-E + zh_Hans: DALL-E 绘画 + pt_BR: DALL-E + description: + en_US: DALL-E art + zh_Hans: DALL-E 绘画 + pt_BR: DALL-E art + icon: icon.png + tags: + - image + - productivity +credentials_for_provider: + openai_api_key: + type: secret-input + required: true + label: + en_US: OpenAI API key + zh_Hans: OpenAI API key + pt_BR: OpenAI API key + help: + en_US: Please input your OpenAI API key + zh_Hans: 请输入你的 OpenAI API key + pt_BR: Please input your OpenAI API key + placeholder: + en_US: Please input your OpenAI API key + zh_Hans: 请输入你的 OpenAI API key + pt_BR: Please input your OpenAI API key + openai_organization_id: + type: text-input + required: false + label: + en_US: OpenAI organization ID + zh_Hans: OpenAI organization ID + pt_BR: OpenAI organization ID + help: + en_US: Please input your OpenAI organization ID + zh_Hans: 请输入你的 OpenAI organization ID + pt_BR: Please input your OpenAI organization ID + placeholder: + en_US: Please input your OpenAI organization ID + zh_Hans: 请输入你的 OpenAI organization ID + pt_BR: Please input your OpenAI organization ID + openai_base_url: + type: text-input + required: false + label: + en_US: OpenAI base URL + zh_Hans: OpenAI base URL + pt_BR: OpenAI base URL + help: + en_US: Please input your OpenAI base URL + zh_Hans: 请输入你的 OpenAI base URL + pt_BR: Please input your OpenAI base URL + placeholder: + en_US: Please input your OpenAI base URL + zh_Hans: 请输入你的 OpenAI base URL + pt_BR: Please input your OpenAI base URL diff --git a/api/core/tools/provider/builtin/dalle/tools/dalle2.py b/api/core/tools/provider/builtin/dalle/tools/dalle2.py new file mode 100644 index 0000000000000000000000000000000000000000..fbd7397292155e64dc909ae3b7b686675f9a5e31 --- /dev/null +++ b/api/core/tools/provider/builtin/dalle/tools/dalle2.py @@ -0,0 +1,66 @@ +from base64 import b64decode +from typing import Any, Union + +from openai import OpenAI +from yarl import URL + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class DallE2Tool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + openai_organization = self.runtime.credentials.get("openai_organization_id", None) + if not openai_organization: + openai_organization = None + openai_base_url = self.runtime.credentials.get("openai_base_url", None) + if not openai_base_url: + openai_base_url = None + else: + openai_base_url = str(URL(openai_base_url) / "v1") + + client = OpenAI( + api_key=self.runtime.credentials["openai_api_key"], + base_url=openai_base_url, + organization=openai_organization, + ) + + SIZE_MAPPING = { + "small": "256x256", + "medium": "512x512", + "large": "1024x1024", + } + + # prompt + prompt = tool_parameters.get("prompt", "") + if not prompt: + return self.create_text_message("Please input prompt") + + # get size + size = SIZE_MAPPING[tool_parameters.get("size", "large")] + + # get n + n = tool_parameters.get("n", 1) + + # call openapi dalle2 + response = client.images.generate(prompt=prompt, model="dall-e-2", size=size, n=n, response_format="b64_json") + + result = [] + + for image in response.data: + result.append( + self.create_blob_message( + blob=b64decode(image.b64_json), + meta={"mime_type": "image/png"}, + save_as=self.VariableKey.IMAGE.value, + ) + ) + + return result diff --git a/api/core/tools/provider/builtin/dalle/tools/dalle2.yaml b/api/core/tools/provider/builtin/dalle/tools/dalle2.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e43e5df8cddd9b8223d6381033e92bd81ce4e773 --- /dev/null +++ b/api/core/tools/provider/builtin/dalle/tools/dalle2.yaml @@ -0,0 +1,74 @@ +identity: + name: dalle2 + author: Dify + label: + en_US: DALL-E 2 + zh_Hans: DALL-E 2 绘画 + description: + en_US: DALL-E 2 is a powerful drawing tool that can draw the image you want based on your prompt + zh_Hans: DALL-E 2 是一个强大的绘画工具,它可以根据您的提示词绘制出您想要的图像 + pt_BR: DALL-E 2 is a powerful drawing tool that can draw the image you want based on your prompt +description: + human: + en_US: DALL-E is a text to image tool + zh_Hans: DALL-E 是一个文本到图像的工具 + pt_BR: DALL-E is a text to image tool + llm: DALL-E is a tool used to generate images from text +parameters: + - name: prompt + type: string + required: true + label: + en_US: Prompt + zh_Hans: 提示词 + pt_BR: Prompt + human_description: + en_US: Image prompt, you can check the official documentation of DallE 2 + zh_Hans: 图像提示词,您可以查看 DallE 2 的官方文档 + pt_BR: Image prompt, you can check the official documentation of DallE 2 + llm_description: Image prompt of DallE 2, you should describe the image you want to generate as a list of words as possible as detailed + form: llm + - name: size + type: select + required: true + human_description: + en_US: used for selecting the image size + zh_Hans: 用于选择图像大小 + pt_BR: used for selecting the image size + label: + en_US: Image size + zh_Hans: 图像大小 + pt_BR: Image size + form: form + options: + - value: small + label: + en_US: Small(256x256) + zh_Hans: 小(256x256) + pt_BR: Small(256x256) + - value: medium + label: + en_US: Medium(512x512) + zh_Hans: 中(512x512) + pt_BR: Medium(512x512) + - value: large + label: + en_US: Large(1024x1024) + zh_Hans: 大(1024x1024) + pt_BR: Large(1024x1024) + default: large + - name: n + type: number + required: true + human_description: + en_US: used for selecting the number of images + zh_Hans: 用于选择图像数量 + pt_BR: used for selecting the number of images + label: + en_US: Number of images + zh_Hans: 图像数量 + pt_BR: Number of images + form: form + default: 1 + min: 1 + max: 10 diff --git a/api/core/tools/provider/builtin/dalle/tools/dalle3.py b/api/core/tools/provider/builtin/dalle/tools/dalle3.py new file mode 100644 index 0000000000000000000000000000000000000000..af9aa6abb4bc3d9ae6e1f0cba243d373c1c4965e --- /dev/null +++ b/api/core/tools/provider/builtin/dalle/tools/dalle3.py @@ -0,0 +1,115 @@ +import base64 +import random +from typing import Any, Union + +from openai import OpenAI +from yarl import URL + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class DallE3Tool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + openai_organization = self.runtime.credentials.get("openai_organization_id", None) + if not openai_organization: + openai_organization = None + openai_base_url = self.runtime.credentials.get("openai_base_url", None) + if not openai_base_url: + openai_base_url = None + else: + openai_base_url = str(URL(openai_base_url) / "v1") + + client = OpenAI( + api_key=self.runtime.credentials["openai_api_key"], + base_url=openai_base_url, + organization=openai_organization, + ) + + SIZE_MAPPING = { + "square": "1024x1024", + "vertical": "1024x1792", + "horizontal": "1792x1024", + } + + # prompt + prompt = tool_parameters.get("prompt", "") + if not prompt: + return self.create_text_message("Please input prompt") + # get size + size = SIZE_MAPPING[tool_parameters.get("size", "square")] + # get n + n = tool_parameters.get("n", 1) + # get quality + quality = tool_parameters.get("quality", "standard") + if quality not in {"standard", "hd"}: + return self.create_text_message("Invalid quality") + # get style + style = tool_parameters.get("style", "vivid") + if style not in {"natural", "vivid"}: + return self.create_text_message("Invalid style") + + # call openapi dalle3 + response = client.images.generate( + prompt=prompt, model="dall-e-3", size=size, n=n, style=style, quality=quality, response_format="b64_json" + ) + + result = [] + + for image in response.data: + mime_type, blob_image = DallE3Tool._decode_image(image.b64_json) + blob_message = self.create_blob_message( + blob=blob_image, meta={"mime_type": mime_type}, save_as=self.VariableKey.IMAGE + ) + result.append(blob_message) + return result + + @staticmethod + def _decode_image(base64_image: str) -> tuple[str, bytes]: + """ + Decode a base64 encoded image. If the image is not prefixed with a MIME type, + it assumes 'image/png' as the default. + + :param base64_image: Base64 encoded image string + :return: A tuple containing the MIME type and the decoded image bytes + """ + if DallE3Tool._is_plain_base64(base64_image): + return "image/png", base64.b64decode(base64_image) + else: + return DallE3Tool._extract_mime_and_data(base64_image) + + @staticmethod + def _is_plain_base64(encoded_str: str) -> bool: + """ + Check if the given encoded string is plain base64 without a MIME type prefix. + + :param encoded_str: Base64 encoded image string + :return: True if the string is plain base64, False otherwise + """ + return not encoded_str.startswith("data:image") + + @staticmethod + def _extract_mime_and_data(encoded_str: str) -> tuple[str, bytes]: + """ + Extract MIME type and image data from a base64 encoded string with a MIME type prefix. + + :param encoded_str: Base64 encoded image string with MIME type prefix + :return: A tuple containing the MIME type and the decoded image bytes + """ + mime_type = encoded_str.split(";")[0].split(":")[1] + image_data_base64 = encoded_str.split(",")[1] + decoded_data = base64.b64decode(image_data_base64) + return mime_type, decoded_data + + @staticmethod + def _generate_random_id(length=8): + characters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789" + random_id = "".join(random.choices(characters, k=length)) + return random_id diff --git a/api/core/tools/provider/builtin/dalle/tools/dalle3.yaml b/api/core/tools/provider/builtin/dalle/tools/dalle3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0cea8af761e1e567102db93f54649cd782caad4e --- /dev/null +++ b/api/core/tools/provider/builtin/dalle/tools/dalle3.yaml @@ -0,0 +1,123 @@ +identity: + name: dalle3 + author: Dify + label: + en_US: DALL-E 3 + zh_Hans: DALL-E 3 绘画 + pt_BR: DALL-E 3 + description: + en_US: DALL-E 3 is a powerful drawing tool that can draw the image you want based on your prompt, compared to DallE 2, DallE 3 has stronger drawing ability, but it will consume more resources + zh_Hans: DALL-E 3 是一个强大的绘画工具,它可以根据您的提示词绘制出您想要的图像,相比于DallE 2, DallE 3拥有更强的绘画能力,但会消耗更多的资源 + pt_BR: DALL-E 3 is a powerful drawing tool that can draw the image you want based on your prompt, compared to DallE 2, DallE 3 has stronger drawing ability, but it will consume more resources +description: + human: + en_US: DALL-E is a text to image tool + zh_Hans: DALL-E 是一个文本到图像的工具 + pt_BR: DALL-E is a text to image tool + llm: DALL-E is a tool used to generate images from text +parameters: + - name: prompt + type: string + required: true + label: + en_US: Prompt + zh_Hans: 提示词 + pt_BR: Prompt + human_description: + en_US: Image prompt, you can check the official documentation of DallE 3 + zh_Hans: 图像提示词,您可以查看 DallE 3 的官方文档 + pt_BR: Image prompt, you can check the official documentation of DallE 3 + llm_description: Image prompt of DallE 3, you should describe the image you want to generate as a list of words as possible as detailed + form: llm + - name: size + type: select + required: true + human_description: + en_US: selecting the image size + zh_Hans: 选择图像大小 + pt_BR: selecting the image size + label: + en_US: Image size + zh_Hans: 图像大小 + pt_BR: Image size + form: form + options: + - value: square + label: + en_US: Squre(1024x1024) + zh_Hans: 方(1024x1024) + pt_BR: Squre(1024x1024) + - value: vertical + label: + en_US: Vertical(1024x1792) + zh_Hans: 竖屏(1024x1792) + pt_BR: Vertical(1024x1792) + - value: horizontal + label: + en_US: Horizontal(1792x1024) + zh_Hans: 横屏(1792x1024) + pt_BR: Horizontal(1792x1024) + default: square + - name: n + type: number + required: true + human_description: + en_US: selecting the number of images + zh_Hans: 选择图像数量 + pt_BR: selecting the number of images + label: + en_US: Number of images + zh_Hans: 图像数量 + pt_BR: Number of images + form: form + min: 1 + max: 1 + default: 1 + - name: quality + type: select + required: true + human_description: + en_US: selecting the image quality + zh_Hans: 选择图像质量 + pt_BR: selecting the image quality + label: + en_US: Image quality + zh_Hans: 图像质量 + pt_BR: Image quality + form: form + options: + - value: standard + label: + en_US: Standard + zh_Hans: 标准 + pt_BR: Standard + - value: hd + label: + en_US: HD + zh_Hans: 高清 + pt_BR: HD + default: standard + - name: style + type: select + required: true + human_description: + en_US: selecting the image style + zh_Hans: 选择图像风格 + pt_BR: selecting the image style + label: + en_US: Image style + zh_Hans: 图像风格 + pt_BR: Image style + form: form + options: + - value: vivid + label: + en_US: Vivid + zh_Hans: 生动 + pt_BR: Vivid + - value: natural + label: + en_US: Natural + zh_Hans: 自然 + pt_BR: Natural + default: vivid diff --git a/api/core/tools/provider/builtin/devdocs/_assets/icon.svg b/api/core/tools/provider/builtin/devdocs/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..c7a19fabfb18bf03811571796f17b60ad2c87a8c --- /dev/null +++ b/api/core/tools/provider/builtin/devdocs/_assets/icon.svg @@ -0,0 +1,4 @@ + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/devdocs/devdocs.py b/api/core/tools/provider/builtin/devdocs/devdocs.py new file mode 100644 index 0000000000000000000000000000000000000000..446c1e548935c0ec394706f37d653952436105f8 --- /dev/null +++ b/api/core/tools/provider/builtin/devdocs/devdocs.py @@ -0,0 +1,21 @@ +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.devdocs.tools.searchDevDocs import SearchDevDocsTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class DevDocsProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + try: + SearchDevDocsTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "doc": "python~3.12", + "topic": "library/code", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/devdocs/devdocs.yaml b/api/core/tools/provider/builtin/devdocs/devdocs.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7552f5a4973f7e1942f17c3c79dab07249b147a7 --- /dev/null +++ b/api/core/tools/provider/builtin/devdocs/devdocs.yaml @@ -0,0 +1,13 @@ +identity: + author: Richards Tu + name: devdocs + label: + en_US: DevDocs + zh_Hans: DevDocs + description: + en_US: Get official developer documentations on DevDocs. + zh_Hans: 从DevDocs获取官方开发者文档。 + icon: icon.svg + tags: + - search + - productivity diff --git a/api/core/tools/provider/builtin/devdocs/tools/searchDevDocs.py b/api/core/tools/provider/builtin/devdocs/tools/searchDevDocs.py new file mode 100644 index 0000000000000000000000000000000000000000..57cf6d7a308dbac76d522d2eb9b94f5dba954b3c --- /dev/null +++ b/api/core/tools/provider/builtin/devdocs/tools/searchDevDocs.py @@ -0,0 +1,47 @@ +from typing import Any, Union + +import requests +from pydantic import BaseModel, Field + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class SearchDevDocsInput(BaseModel): + doc: str = Field(..., description="The name of the documentation.") + topic: str = Field(..., description="The path of the section/topic.") + + +class SearchDevDocsTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Invokes the DevDocs search tool with the given user ID and tool parameters. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (dict[str, Any]): The parameters for the tool, including 'doc' and 'topic'. + + Returns: + ToolInvokeMessage | list[ToolInvokeMessage]: The result of the tool invocation, + which can be a single message or a list of messages. + """ + doc = tool_parameters.get("doc", "") + topic = tool_parameters.get("topic", "") + + if not doc: + return self.create_text_message("Please provide the documentation name.") + if not topic: + return self.create_text_message("Please provide the topic path.") + + url = f"https://documents.devdocs.io/{doc}/{topic}.html" + response = requests.get(url) + + if response.status_code == 200: + content = response.text + return self.create_text_message(self.summary(user_id=user_id, content=content)) + else: + return self.create_text_message( + f"Failed to retrieve the documentation. Status code: {response.status_code}" + ) diff --git a/api/core/tools/provider/builtin/devdocs/tools/searchDevDocs.yaml b/api/core/tools/provider/builtin/devdocs/tools/searchDevDocs.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2476db9da42d60fc79da642cac9b07205f04a89f --- /dev/null +++ b/api/core/tools/provider/builtin/devdocs/tools/searchDevDocs.yaml @@ -0,0 +1,34 @@ +identity: + name: searchDevDocs + author: Richards Tu + label: + en_US: Search Developer Docs + zh_Hans: 搜索开发者文档 +description: + human: + en_US: A tools for searching for a specific topic and path in DevDocs based on the provided documentation name and topic. Don't for get to add some shots in the system prompt; for example, the documentation name should be like \"vuex~4\", \"css\", or \"python~3.12\", while the topic should be like \"guide/actions\" for Vuex 4, \"display-box\" for CSS, or \"library/code\" for Python 3.12. + zh_Hans: 一个用于根据提供的文档名称和主题,在DevDocs中搜索特定主题和路径的工具。不要忘记在系统提示词中添加一些示例;例如,文档名称应该是\"vuex~4\"、\"css\"或\"python~3.12\",而主题应该是\"guide/actions\"用于Vuex 4,\"display-box\"用于CSS,或\"library/code\"用于Python 3.12。 + llm: A tools for searching for specific developer documentation in DevDocs based on the provided documentation name and topic. +parameters: + - name: doc + type: string + required: true + label: + en_US: Documentation name + zh_Hans: 文档名称 + human_description: + en_US: The name of the documentation. + zh_Hans: 文档名称。 + llm_description: The name of the documentation, such as \"vuex~4\", \"css\", or \"python~3.12\". The exact value should be identified by the user. + form: llm + - name: topic + type: string + required: true + label: + en_US: Topic name + zh_Hans: 主题名称 + human_description: + en_US: The path of the section/topic. + zh_Hans: 文档主题的路径。 + llm_description: The path of the section/topic, such as \"guide/actions\" for Vuex 4, \"display-box\" for CSS, or \"library/code\" for Python 3.12. + form: llm diff --git a/api/core/tools/provider/builtin/did/_assets/icon.svg b/api/core/tools/provider/builtin/did/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..c477d7cb71dea28ec0fbd79dfd8ec6194244f26d --- /dev/null +++ b/api/core/tools/provider/builtin/did/_assets/icon.svg @@ -0,0 +1,14 @@ + + + + + + + + + + + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/did/did.py b/api/core/tools/provider/builtin/did/did.py new file mode 100644 index 0000000000000000000000000000000000000000..5af78794f625b7b33ecdefc89e042e3057ad535d --- /dev/null +++ b/api/core/tools/provider/builtin/did/did.py @@ -0,0 +1,18 @@ +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.did.tools.talks import TalksTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class DIDProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + try: + # Example validation using the D-ID talks tool + TalksTool().fork_tool_runtime(runtime={"credentials": credentials}).invoke( + user_id="", + tool_parameters={ + "source_url": "https://www.d-id.com/wp-content/uploads/2023/11/Hero-image-1.png", + "text_input": "Hello, welcome to use D-ID tool in Dify", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/did/did.yaml b/api/core/tools/provider/builtin/did/did.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a70b71812e46485a3489daba50c2d9af6e2da373 --- /dev/null +++ b/api/core/tools/provider/builtin/did/did.yaml @@ -0,0 +1,28 @@ +identity: + author: Matri Qi + name: did + label: + en_US: D-ID + description: + en_US: D-ID is a tool enabling the creation of high-quality, custom videos of Digital Humans from a single image. + icon: icon.svg + tags: + - videos +credentials_for_provider: + did_api_key: + type: secret-input + required: true + label: + en_US: D-ID API Key + placeholder: + en_US: Please input your D-ID API key + help: + en_US: Get your D-ID API key from your D-ID account settings. + url: https://studio.d-id.com/account-settings + base_url: + type: text-input + required: false + label: + en_US: D-ID server's Base URL + placeholder: + en_US: https://api.d-id.com diff --git a/api/core/tools/provider/builtin/did/did_appx.py b/api/core/tools/provider/builtin/did/did_appx.py new file mode 100644 index 0000000000000000000000000000000000000000..dca62f9e198262612a84273b232f04a5107f749d --- /dev/null +++ b/api/core/tools/provider/builtin/did/did_appx.py @@ -0,0 +1,87 @@ +import logging +import time +from collections.abc import Mapping +from typing import Any + +import requests +from requests.exceptions import HTTPError + +logger = logging.getLogger(__name__) + + +class DIDApp: + def __init__(self, api_key: str | None = None, base_url: str | None = None): + self.api_key = api_key + self.base_url = base_url or "https://api.d-id.com" + if not self.api_key: + raise ValueError("API key is required") + + def _prepare_headers(self, idempotency_key: str | None = None): + headers = {"Content-Type": "application/json", "Authorization": f"Basic {self.api_key}"} + if idempotency_key: + headers["Idempotency-Key"] = idempotency_key + return headers + + def _request( + self, + method: str, + url: str, + data: Mapping[str, Any] | None = None, + headers: Mapping[str, str] | None = None, + retries: int = 3, + backoff_factor: float = 0.3, + ) -> Mapping[str, Any] | None: + for i in range(retries): + try: + response = requests.request(method, url, json=data, headers=headers) + response.raise_for_status() + return response.json() + except requests.exceptions.RequestException as e: + if i < retries - 1 and isinstance(e, HTTPError) and e.response.status_code >= 500: + time.sleep(backoff_factor * (2**i)) + else: + raise + return None + + def talks(self, wait: bool = True, poll_interval: int = 5, idempotency_key: str | None = None, **kwargs): + endpoint = f"{self.base_url}/talks" + headers = self._prepare_headers(idempotency_key) + data = kwargs["params"] + logger.debug(f"Send request to {endpoint=} body={data}") + response = self._request("POST", endpoint, data, headers) + if response is None: + raise HTTPError("Failed to initiate D-ID talks after multiple retries") + id: str = response["id"] + if wait: + return self._monitor_job_status(id=id, target="talks", poll_interval=poll_interval) + return id + + def animations(self, wait: bool = True, poll_interval: int = 5, idempotency_key: str | None = None, **kwargs): + endpoint = f"{self.base_url}/animations" + headers = self._prepare_headers(idempotency_key) + data = kwargs["params"] + logger.debug(f"Send request to {endpoint=} body={data}") + response = self._request("POST", endpoint, data, headers) + if response is None: + raise HTTPError("Failed to initiate D-ID talks after multiple retries") + id: str = response["id"] + if wait: + return self._monitor_job_status(target="animations", id=id, poll_interval=poll_interval) + return id + + def check_did_status(self, target: str, id: str): + endpoint = f"{self.base_url}/{target}/{id}" + headers = self._prepare_headers() + response = self._request("GET", endpoint, headers=headers) + if response is None: + raise HTTPError(f"Failed to check status for talks {id} after multiple retries") + return response + + def _monitor_job_status(self, target: str, id: str, poll_interval: int): + while True: + status = self.check_did_status(target=target, id=id) + if status["status"] == "done": + return status + elif status["status"] == "error" or status["status"] == "rejected": + raise HTTPError(f"Talks {id} failed: {status['status']} {status.get('error', {}).get('description')}") + time.sleep(poll_interval) diff --git a/api/core/tools/provider/builtin/did/tools/animations.py b/api/core/tools/provider/builtin/did/tools/animations.py new file mode 100644 index 0000000000000000000000000000000000000000..bc9d17e40d2878cdcae3cc2d6d15beacf20644df --- /dev/null +++ b/api/core/tools/provider/builtin/did/tools/animations.py @@ -0,0 +1,49 @@ +import json +from typing import Any, Union + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.did.did_appx import DIDApp +from core.tools.tool.builtin_tool import BuiltinTool + + +class AnimationsTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + app = DIDApp(api_key=self.runtime.credentials["did_api_key"], base_url=self.runtime.credentials["base_url"]) + + driver_expressions_str = tool_parameters.get("driver_expressions") + driver_expressions = json.loads(driver_expressions_str) if driver_expressions_str else None + + config = { + "stitch": tool_parameters.get("stitch", True), + "mute": tool_parameters.get("mute"), + "result_format": tool_parameters.get("result_format") or "mp4", + } + config = {k: v for k, v in config.items() if v is not None and v != ""} + + options = { + "source_url": tool_parameters["source_url"], + "driver_url": tool_parameters.get("driver_url"), + "config": config, + } + options = {k: v for k, v in options.items() if v is not None and v != ""} + + if not options.get("source_url"): + raise ValueError("Source URL is required") + + if config.get("logo_url"): + if not config.get("logo_x"): + raise ValueError("Logo X position is required when logo URL is provided") + if not config.get("logo_y"): + raise ValueError("Logo Y position is required when logo URL is provided") + + animations_result = app.animations(params=options, wait=True) + + if not isinstance(animations_result, str): + animations_result = json.dumps(animations_result, ensure_ascii=False, indent=4) + + if not animations_result: + return self.create_text_message("D-ID animations request failed.") + + return self.create_text_message(animations_result) diff --git a/api/core/tools/provider/builtin/did/tools/animations.yaml b/api/core/tools/provider/builtin/did/tools/animations.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2a2036c7b2a88fc10ec3d476c1b66ec15a1d39d5 --- /dev/null +++ b/api/core/tools/provider/builtin/did/tools/animations.yaml @@ -0,0 +1,86 @@ +identity: + name: animations + author: Matri Qi + label: + en_US: Animations +description: + human: + en_US: Animations enables to create videos matching head movements, expressions, emotions, and voice from a driver video and image. + llm: Animations enables to create videos matching head movements, expressions, emotions, and voice from a driver video and image. +parameters: + - name: source_url + type: string + required: true + label: + en_US: source url + human_description: + en_US: The URL of the source image to be animated by the driver video, or a selection from the list of provided studio actors. + llm_description: The URL of the source image to be animated by the driver video, or a selection from the list of provided studio actors. + form: llm + - name: driver_url + type: string + required: false + label: + en_US: driver url + human_description: + en_US: The URL of the driver video to drive the animation, or a provided driver name from D-ID. + form: form + - name: mute + type: boolean + required: false + label: + en_US: mute + human_description: + en_US: Mutes the driver sound in the animated video result, defaults to true + form: form + - name: stitch + type: boolean + required: false + label: + en_US: stitch + human_description: + en_US: If enabled, the driver video will be stitched with the animationing head video. + form: form + - name: logo_url + type: string + required: false + label: + en_US: logo url + human_description: + en_US: The URL of the logo image to be added to the animation video. + form: form + - name: logo_x + type: number + required: false + label: + en_US: logo position x + human_description: + en_US: The x position of the logo image in the animation video. It's required when logo url is provided. + form: form + - name: logo_y + type: number + required: false + label: + en_US: logo position y + human_description: + en_US: The y position of the logo image in the animation video. It's required when logo url is provided. + form: form + - name: result_format + type: string + default: mp4 + required: false + label: + en_US: result format + human_description: + en_US: The format of the result video. + form: form + options: + - value: mp4 + label: + en_US: mp4 + - value: gif + label: + en_US: gif + - value: mov + label: + en_US: mov diff --git a/api/core/tools/provider/builtin/did/tools/talks.py b/api/core/tools/provider/builtin/did/tools/talks.py new file mode 100644 index 0000000000000000000000000000000000000000..d6f0c7ff1797932b102f42bd1543462ed1db5e65 --- /dev/null +++ b/api/core/tools/provider/builtin/did/tools/talks.py @@ -0,0 +1,65 @@ +import json +from typing import Any, Union + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.did.did_appx import DIDApp +from core.tools.tool.builtin_tool import BuiltinTool + + +class TalksTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + app = DIDApp(api_key=self.runtime.credentials["did_api_key"], base_url=self.runtime.credentials["base_url"]) + + driver_expressions_str = tool_parameters.get("driver_expressions") + driver_expressions = json.loads(driver_expressions_str) if driver_expressions_str else None + + script = { + "type": tool_parameters.get("script_type") or "text", + "input": tool_parameters.get("text_input"), + "audio_url": tool_parameters.get("audio_url"), + "reduce_noise": tool_parameters.get("audio_reduce_noise", False), + } + script = {k: v for k, v in script.items() if v is not None and v != ""} + config = { + "stitch": tool_parameters.get("stitch", True), + "sharpen": tool_parameters.get("sharpen"), + "fluent": tool_parameters.get("fluent"), + "result_format": tool_parameters.get("result_format") or "mp4", + "pad_audio": tool_parameters.get("pad_audio"), + "driver_expressions": driver_expressions, + } + config = {k: v for k, v in config.items() if v is not None and v != ""} + + options = { + "source_url": tool_parameters["source_url"], + "driver_url": tool_parameters.get("driver_url"), + "script": script, + "config": config, + } + options = {k: v for k, v in options.items() if v is not None and v != ""} + + if not options.get("source_url"): + raise ValueError("Source URL is required") + + if script.get("type") == "audio": + script.pop("input", None) + if not script.get("audio_url"): + raise ValueError("Audio URL is required for audio script type") + + if script.get("type") == "text": + script.pop("audio_url", None) + script.pop("reduce_noise", None) + if not script.get("input"): + raise ValueError("Text input is required for text script type") + + talks_result = app.talks(params=options, wait=True) + + if not isinstance(talks_result, str): + talks_result = json.dumps(talks_result, ensure_ascii=False, indent=4) + + if not talks_result: + return self.create_text_message("D-ID talks request failed.") + + return self.create_text_message(talks_result) diff --git a/api/core/tools/provider/builtin/did/tools/talks.yaml b/api/core/tools/provider/builtin/did/tools/talks.yaml new file mode 100644 index 0000000000000000000000000000000000000000..88d430512923e4337bfd961c9b946cdf199a05a4 --- /dev/null +++ b/api/core/tools/provider/builtin/did/tools/talks.yaml @@ -0,0 +1,126 @@ +identity: + name: talks + author: Matri Qi + label: + en_US: Talks +description: + human: + en_US: Talks enables the creation of realistic talking head videos from text or audio inputs. + llm: Talks enables the creation of realistic talking head videos from text or audio inputs. +parameters: + - name: source_url + type: string + required: true + label: + en_US: source url + human_description: + en_US: The URL of the source image to be animated by the driver video, or a selection from the list of provided studio actors. + llm_description: The URL of the source image to be animated by the driver video, or a selection from the list of provided studio actors. + form: llm + - name: driver_url + type: string + required: false + label: + en_US: driver url + human_description: + en_US: The URL of the driver video to drive the talk, or a provided driver name from D-ID. + form: form + - name: script_type + type: string + required: false + label: + en_US: script type + human_description: + en_US: The type of the script. + form: form + options: + - value: text + label: + en_US: text + - value: audio + label: + en_US: audio + - name: text_input + type: string + required: false + label: + en_US: text input + human_description: + en_US: The text input to be spoken by the talking head. Required when script type is text. + form: form + - name: audio_url + type: string + required: false + label: + en_US: audio url + human_description: + en_US: The URL of the audio file to be spoken by the talking head. Required when script type is audio. + form: form + - name: audio_reduce_noise + type: boolean + required: false + label: + en_US: audio reduce noise + human_description: + en_US: If enabled, the audio will be processed to reduce noise before being spoken by the talking head. It only works when script type is audio. + form: form + - name: stitch + type: boolean + required: false + label: + en_US: stitch + human_description: + en_US: If enabled, the driver video will be stitched with the talking head video. + form: form + - name: sharpen + type: boolean + required: false + label: + en_US: sharpen + human_description: + en_US: If enabled, the talking head video will be sharpened. + form: form + - name: result_format + type: string + required: false + label: + en_US: result format + human_description: + en_US: The format of the result video. + form: form + options: + - value: mp4 + label: + en_US: mp4 + - value: gif + label: + en_US: gif + - value: mov + label: + en_US: mov + - name: fluent + type: boolean + required: false + label: + en_US: fluent + human_description: + en_US: Interpolate between the last & first frames of the driver video When used together with pad_audio can create a seamless transition between videos of the same driver + form: form + - name: pad_audio + type: number + required: false + label: + en_US: pad audio + human_description: + en_US: Pad the audio with silence at the end (given in seconds) Will increase the video duration & the credits it consumes + form: form + min: 1 + max: 60 + - name: driver_expressions + type: string + required: false + label: + en_US: driver expressions + human_description: + en_US: timed expressions for animation. It should be an JSON array style string. Take D-ID documentation(https://docs.d-id.com/reference/createtalk) for more information. + form: form diff --git a/api/core/tools/provider/builtin/dingtalk/_assets/icon.svg b/api/core/tools/provider/builtin/dingtalk/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..b60653b7a59409798865a3046a0139f42995d80c --- /dev/null +++ b/api/core/tools/provider/builtin/dingtalk/_assets/icon.svg @@ -0,0 +1,7 @@ + + + + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/dingtalk/dingtalk.py b/api/core/tools/provider/builtin/dingtalk/dingtalk.py new file mode 100644 index 0000000000000000000000000000000000000000..be1d5e099c22462941dbc19b8c20f8a4122597a2 --- /dev/null +++ b/api/core/tools/provider/builtin/dingtalk/dingtalk.py @@ -0,0 +1,8 @@ +from core.tools.provider.builtin.dingtalk.tools.dingtalk_group_bot import DingTalkGroupBotTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class DingTalkProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + DingTalkGroupBotTool() + pass diff --git a/api/core/tools/provider/builtin/dingtalk/dingtalk.yaml b/api/core/tools/provider/builtin/dingtalk/dingtalk.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c922c140a8badc4c7d53051f698f55ad66ac49ad --- /dev/null +++ b/api/core/tools/provider/builtin/dingtalk/dingtalk.yaml @@ -0,0 +1,16 @@ +identity: + author: Bowen Liang + name: dingtalk + label: + en_US: DingTalk + zh_Hans: 钉钉 + pt_BR: DingTalk + description: + en_US: DingTalk group robot + zh_Hans: 钉钉群机器人 + pt_BR: DingTalk group robot + icon: icon.svg + tags: + - social + - productivity +credentials_for_provider: diff --git a/api/core/tools/provider/builtin/dingtalk/tools/dingtalk_group_bot.py b/api/core/tools/provider/builtin/dingtalk/tools/dingtalk_group_bot.py new file mode 100644 index 0000000000000000000000000000000000000000..f33ad5be59b4031f5623bddf2b6337f2563651cb --- /dev/null +++ b/api/core/tools/provider/builtin/dingtalk/tools/dingtalk_group_bot.py @@ -0,0 +1,89 @@ +import base64 +import hashlib +import hmac +import logging +import time +import urllib.parse +from typing import Any, Union + +import httpx + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class DingTalkGroupBotTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + Dingtalk custom group robot API docs: + https://open.dingtalk.com/document/orgapp/custom-robot-access + """ + content = tool_parameters.get("content") + if not content: + return self.create_text_message("Invalid parameter content") + + access_token = tool_parameters.get("access_token") + if not access_token: + return self.create_text_message( + "Invalid parameter access_token. " + "Regarding information about security details," + "please refer to the DingTalk docs:" + "https://open.dingtalk.com/document/robots/customize-robot-security-settings" + ) + + sign_secret = tool_parameters.get("sign_secret") + if not sign_secret: + return self.create_text_message( + "Invalid parameter sign_secret. " + "Regarding information about security details," + "please refer to the DingTalk docs:" + "https://open.dingtalk.com/document/robots/customize-robot-security-settings" + ) + + msgtype = "text" + api_url = "https://oapi.dingtalk.com/robot/send" + headers = { + "Content-Type": "application/json", + } + params = { + "access_token": access_token, + } + + self._apply_security_mechanism(params, sign_secret) + + payload = { + "msgtype": msgtype, + "text": { + "content": content, + }, + } + + try: + res = httpx.post(api_url, headers=headers, params=params, json=payload) + if res.is_success: + return self.create_text_message("Text message sent successfully") + else: + return self.create_text_message( + f"Failed to send the text message, status code: {res.status_code}, response: {res.text}" + ) + except Exception as e: + return self.create_text_message("Failed to send message to group chat bot. {}".format(e)) + + @staticmethod + def _apply_security_mechanism(params: dict[str, Any], sign_secret: str): + try: + timestamp = str(round(time.time() * 1000)) + secret_enc = sign_secret.encode("utf-8") + string_to_sign = f"{timestamp}\n{sign_secret}" + string_to_sign_enc = string_to_sign.encode("utf-8") + hmac_code = hmac.new(secret_enc, string_to_sign_enc, digestmod=hashlib.sha256).digest() + sign = urllib.parse.quote_plus(base64.b64encode(hmac_code)) + + params["timestamp"] = timestamp + params["sign"] = sign + except Exception: + msg = "Failed to apply security mechanism to the request." + logging.exception(msg) diff --git a/api/core/tools/provider/builtin/dingtalk/tools/dingtalk_group_bot.yaml b/api/core/tools/provider/builtin/dingtalk/tools/dingtalk_group_bot.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dc8a90b71939032bbe44367f502964759ecba2c7 --- /dev/null +++ b/api/core/tools/provider/builtin/dingtalk/tools/dingtalk_group_bot.yaml @@ -0,0 +1,52 @@ +identity: + name: dingtalk_group_bot + author: Bowen Liang + label: + en_US: Send Group Message + zh_Hans: 发送群消息 + pt_BR: Send Group Message + icon: icon.svg +description: + human: + en_US: Sending a group message on DingTalk via the webhook of group bot + zh_Hans: 通过钉钉的群机器人webhook发送群消息 + pt_BR: Sending a group message on DingTalk via the webhook of group bot + llm: A tool for sending messages to a chat group on DingTalk(钉钉) . +parameters: + - name: access_token + type: secret-input + required: true + label: + en_US: access token + zh_Hans: access token + pt_BR: access token + human_description: + en_US: access_token in the group robot webhook + zh_Hans: 群自定义机器人webhook中access_token字段的值 + pt_BR: access_token in the group robot webhook + form: form + - name: sign_secret + type: secret-input + required: true + label: + en_US: secret key for signing + zh_Hans: 加签秘钥 + pt_BR: secret key for signing + human_description: + en_US: secret key for signing + zh_Hans: 加签秘钥 + pt_BR: secret key for signing + form: form + - name: content + type: string + required: true + label: + en_US: content + zh_Hans: 消息内容 + pt_BR: content + human_description: + en_US: Content to sent to the group. + zh_Hans: 群消息文本 + pt_BR: Content to sent to the group. + llm_description: Content of the message + form: llm diff --git a/api/core/tools/provider/builtin/discord/_assets/icon.svg b/api/core/tools/provider/builtin/discord/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..177a0591f9cb08738c3b8abad743bd92ae329cc1 --- /dev/null +++ b/api/core/tools/provider/builtin/discord/_assets/icon.svg @@ -0,0 +1,7 @@ + + + + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/discord/discord.py b/api/core/tools/provider/builtin/discord/discord.py new file mode 100644 index 0000000000000000000000000000000000000000..c94824b591cd95eefde8ee5089065edbca59726f --- /dev/null +++ b/api/core/tools/provider/builtin/discord/discord.py @@ -0,0 +1,9 @@ +from typing import Any + +from core.tools.provider.builtin.discord.tools.discord_webhook import DiscordWebhookTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class DiscordProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + DiscordWebhookTool() diff --git a/api/core/tools/provider/builtin/discord/discord.yaml b/api/core/tools/provider/builtin/discord/discord.yaml new file mode 100644 index 0000000000000000000000000000000000000000..18b249b5229a0e81eb087e07efe02d319efbe186 --- /dev/null +++ b/api/core/tools/provider/builtin/discord/discord.yaml @@ -0,0 +1,16 @@ +identity: + author: Ice Yao + name: discord + label: + en_US: Discord + zh_Hans: Discord + pt_BR: Discord + description: + en_US: Discord Webhook + zh_Hans: Discord Webhook + pt_BR: Discord Webhook + icon: icon.svg + tags: + - social + - productivity +credentials_for_provider: diff --git a/api/core/tools/provider/builtin/discord/tools/discord_webhook.py b/api/core/tools/provider/builtin/discord/tools/discord_webhook.py new file mode 100644 index 0000000000000000000000000000000000000000..c1834a1a265be2d9a257ce3617b578d34c124d6a --- /dev/null +++ b/api/core/tools/provider/builtin/discord/tools/discord_webhook.py @@ -0,0 +1,49 @@ +from typing import Any, Union + +import httpx + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class DiscordWebhookTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Incoming Webhooks + API Document: + https://discord.com/developers/docs/resources/webhook#execute-webhook + """ + + content = tool_parameters.get("content", "") + if not content: + return self.create_text_message("Invalid parameter content") + + webhook_url = tool_parameters.get("webhook_url", "") + if not webhook_url.startswith("https://discord.com/api/webhooks/"): + return self.create_text_message( + f"Invalid parameter webhook_url ${webhook_url}, \ + not a valid Discord webhook URL" + ) + + headers = { + "Content-Type": "application/json", + } + payload = { + "username": tool_parameters.get("username") or user_id, + "content": content, + "avatar_url": tool_parameters.get("avatar_url") or None, + } + + try: + res = httpx.post(webhook_url, headers=headers, json=payload) + if res.is_success: + return self.create_text_message("Text message was sent successfully") + else: + return self.create_text_message( + f"Failed to send the text message, \ + status code: {res.status_code}, response: {res.text}" + ) + except Exception as e: + return self.create_text_message("Failed to send message through webhook. {}".format(e)) diff --git a/api/core/tools/provider/builtin/discord/tools/discord_webhook.yaml b/api/core/tools/provider/builtin/discord/tools/discord_webhook.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6847b973cabd195f81f2f14e245dadb79c95d569 --- /dev/null +++ b/api/core/tools/provider/builtin/discord/tools/discord_webhook.yaml @@ -0,0 +1,65 @@ +identity: + name: discord_webhook + author: Ice Yao + label: + en_US: Incoming Webhook to send message + zh_Hans: 通过入站Webhook发送消息 + pt_BR: Incoming Webhook to send message + icon: icon.svg +description: + human: + en_US: Sending a message on Discord via the Incoming Webhook + zh_Hans: 通过入站Webhook在Discord上发送消息 + pt_BR: Sending a message on Discord via the Incoming Webhook + llm: A tool for sending messages to a chat on Discord. +parameters: + - name: webhook_url + type: string + required: true + label: + en_US: Discord Incoming Webhook url + zh_Hans: Discord入站Webhook的url + pt_BR: Discord Incoming Webhook url + human_description: + en_US: Discord Incoming Webhook url + zh_Hans: Discord入站Webhook的url + pt_BR: Discord Incoming Webhook url + form: form + - name: content + type: string + required: true + label: + en_US: content + zh_Hans: 消息内容 + pt_BR: content + human_description: + en_US: Content to sent to the channel or person. + zh_Hans: 消息内容文本 + pt_BR: Content to sent to the channel or person. + llm_description: Content of the message + form: llm + - name: username + type: string + required: false + label: + en_US: Discord Webhook Username + zh_Hans: Discord Webhook用户名 + pt_BR: Discord Webhook Username + human_description: + en_US: Discord Webhook Username + zh_Hans: Discord Webhook用户名 + pt_BR: Discord Webhook Username + llm_description: Discord Webhook Username + form: llm + - name: avatar_url + type: string + required: false + label: + en_US: Discord Webhook Avatar + zh_Hans: Discord Webhook头像 + pt_BR: Discord Webhook Avatar + human_description: + en_US: Discord Webhook Avatar URL + zh_Hans: Discord Webhook头像地址 + pt_BR: Discord Webhook Avatar URL + form: form diff --git a/api/core/tools/provider/builtin/duckduckgo/_assets/icon.svg b/api/core/tools/provider/builtin/duckduckgo/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..a816a6b49ebb787d6c6d348152307619c71e80bd --- /dev/null +++ b/api/core/tools/provider/builtin/duckduckgo/_assets/icon.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/duckduckgo/duckduckgo.py b/api/core/tools/provider/builtin/duckduckgo/duckduckgo.py new file mode 100644 index 0000000000000000000000000000000000000000..8269167127b8e5bdb161e3adbf64a0a14420bccd --- /dev/null +++ b/api/core/tools/provider/builtin/duckduckgo/duckduckgo.py @@ -0,0 +1,20 @@ +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.duckduckgo.tools.ddgo_search import DuckDuckGoSearchTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class DuckDuckGoProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + try: + DuckDuckGoSearchTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "query": "John Doe", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/duckduckgo/duckduckgo.yaml b/api/core/tools/provider/builtin/duckduckgo/duckduckgo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f3faa0604557729175c997bc15d52827a0c41a1b --- /dev/null +++ b/api/core/tools/provider/builtin/duckduckgo/duckduckgo.yaml @@ -0,0 +1,12 @@ +identity: + author: Yash Parmar + name: duckduckgo + label: + en_US: DuckDuckGo + zh_Hans: DuckDuckGo + description: + en_US: A privacy-focused search engine. + zh_Hans: 一个注重隐私的搜索引擎。 + icon: icon.svg + tags: + - search diff --git a/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_ai.py b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_ai.py new file mode 100644 index 0000000000000000000000000000000000000000..8bdd638f4a01d149685692513ca5ce0945a8c860 --- /dev/null +++ b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_ai.py @@ -0,0 +1,20 @@ +from typing import Any + +from duckduckgo_search import DDGS + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class DuckDuckGoAITool(BuiltinTool): + """ + Tool for performing a search using DuckDuckGo search engine. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + query_dict = { + "keywords": tool_parameters.get("query"), + "model": tool_parameters.get("model"), + } + response = DDGS().chat(**query_dict) + return self.create_text_message(text=response) diff --git a/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_ai.yaml b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_ai.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dd049d3b5a13d2c6719a8a65fa825bd81c65d9be --- /dev/null +++ b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_ai.yaml @@ -0,0 +1,47 @@ +identity: + name: ddgo_ai + author: hjlarry + label: + en_US: DuckDuckGo AI Chat + zh_Hans: DuckDuckGo AI聊天 +description: + human: + en_US: Use the anonymous private chat provided by DuckDuckGo. + zh_Hans: 使用DuckDuckGo提供的匿名私密聊天。 + llm: Use the anonymous private chat provided by DuckDuckGo. +parameters: + - name: query + type: string + required: true + label: + en_US: Chat Content + zh_Hans: 聊天内容 + human_description: + en_US: The chat content. + zh_Hans: 要聊天的内容。 + llm_description: Key words for chat + form: llm + - name: model + type: select + required: true + options: + - value: gpt-4o-mini + label: + en_US: GPT-4o-mini + - value: claude-3-haiku + label: + en_US: Claude 3 + - value: llama-3-70b + label: + en_US: Llama 3 + - value: mixtral-8x7b + label: + en_US: Mixtral + default: gpt-4o-mini + label: + en_US: Choose Model + zh_Hans: 选择模型 + human_description: + en_US: used to select the model for AI chat. + zh_Hans: 用于选择使用AI聊天的模型 + form: form diff --git a/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_img.py b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_img.py new file mode 100644 index 0000000000000000000000000000000000000000..b3c630878f3c6230e8d82f0dba3aa9a927d9607a --- /dev/null +++ b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_img.py @@ -0,0 +1,33 @@ +from typing import Any + +from duckduckgo_search import DDGS + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class DuckDuckGoImageSearchTool(BuiltinTool): + """ + Tool for performing an image search using DuckDuckGo search engine. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> list[ToolInvokeMessage]: + query_dict = { + "keywords": tool_parameters.get("query"), + "timelimit": tool_parameters.get("timelimit"), + "size": tool_parameters.get("size"), + "max_results": tool_parameters.get("max_results"), + } + + # Add query_prefix handling + query_prefix = tool_parameters.get("query_prefix", "").strip() + final_query = f"{query_prefix} {query_dict['keywords']}".strip() + query_dict["keywords"] = final_query + + response = DDGS().images(**query_dict) + markdown_result = "\n\n" + json_result = [] + for res in response: + markdown_result += f"![{res.get('title') or ''}]({res.get('image') or ''})" + json_result.append(self.create_json_message(res)) + return [self.create_text_message(markdown_result)] + json_result diff --git a/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_img.yaml b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_img.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a543d1e218b578dc1590fec164b3d92b0b0855a2 --- /dev/null +++ b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_img.yaml @@ -0,0 +1,99 @@ +identity: + name: ddgo_img + author: hjlarry + label: + en_US: DuckDuckGo Image Search + zh_Hans: DuckDuckGo 图片搜索 +description: + human: + en_US: Perform image searches on DuckDuckGo and get results. + zh_Hans: 在 DuckDuckGo 上进行图片搜索并获取结果。 + llm: Perform image searches on DuckDuckGo and get results. +parameters: + - name: query + type: string + required: true + label: + en_US: Query string + zh_Hans: 查询语句 + human_description: + en_US: The search query. + zh_Hans: 搜索查询语句。 + llm_description: Key words for searching + form: llm + - name: max_results + type: number + required: true + default: 3 + label: + en_US: Max results + zh_Hans: 最大结果数量 + human_description: + en_US: The max results. + zh_Hans: 最大结果数量 + form: form + - name: timelimit + type: select + required: false + options: + - value: Day + label: + en_US: current day + zh_Hans: 当天 + - value: Week + label: + en_US: current week + zh_Hans: 本周 + - value: Month + label: + en_US: current month + zh_Hans: 当月 + - value: Year + label: + en_US: current year + zh_Hans: 今年 + label: + en_US: Result time limit + zh_Hans: 结果时间限制 + human_description: + en_US: Use when querying results within a specific time range only. + zh_Hans: 只查询一定时间范围内的结果时使用 + form: form + - name: size + type: select + required: false + options: + - value: Small + label: + en_US: small + zh_Hans: 小 + - value: Medium + label: + en_US: medium + zh_Hans: 中 + - value: Large + label: + en_US: large + zh_Hans: 大 + - value: Wallpaper + label: + en_US: xl + zh_Hans: 超大 + label: + en_US: image size + zh_Hans: 图片大小 + human_description: + en_US: The size of the image to be searched. + zh_Hans: 要搜索的图片的大小 + form: form + - name: query_prefix + label: + en_US: Query Prefix + zh_Hans: 查询前缀 + type: string + required: false + default: "" + form: form + human_description: + en_US: Specific Search e.g. "site:unsplash.com" + zh_Hans: 定向搜索 e.g. "site:unsplash.com" diff --git a/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_news.py b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_news.py new file mode 100644 index 0000000000000000000000000000000000000000..11da6f5cf76580e1f1f398ea69173ecbe5b5d626 --- /dev/null +++ b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_news.py @@ -0,0 +1,93 @@ +from typing import Any + +from duckduckgo_search import DDGS + +from core.model_runtime.entities.message_entities import SystemPromptMessage +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +SUMMARY_PROMPT = """ +User's query: +{query} + +Here are the news results: +{content} + +Please summarize the news in a few sentences. +""" + + +class DuckDuckGoNewsSearchTool(BuiltinTool): + """ + Tool for performing a news search using DuckDuckGo search engine. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage | list[ToolInvokeMessage]: + query_dict = { + "keywords": tool_parameters.get("query"), + "timelimit": tool_parameters.get("timelimit"), + "max_results": tool_parameters.get("max_results"), + "safesearch": "moderate", + "region": "wt-wt", + } + + # Add query_prefix handling + query_prefix = tool_parameters.get("query_prefix", "").strip() + final_query = f"{query_prefix} {query_dict['keywords']}".strip() + query_dict["keywords"] = final_query + + try: + response = list(DDGS().news(**query_dict)) + if not response: + return [self.create_text_message("No news found matching your criteria.")] + except Exception as e: + return [self.create_text_message(f"Error searching news: {str(e)}")] + + require_summary = tool_parameters.get("require_summary", False) + + if require_summary: + results = "\n".join([f"{res.get('title')}: {res.get('body')}" for res in response]) + results = self.summary_results(user_id=user_id, content=results, query=query_dict["keywords"]) + return self.create_text_message(text=results) + + # Create rich markdown content for each news item + markdown_result = "\n\n" + json_result = [] + + for res in response: + markdown_result += f"### {res.get('title', 'Untitled')}\n\n" + if res.get("date"): + markdown_result += f"**Date:** {res.get('date')}\n\n" + if res.get("body"): + markdown_result += f"{res.get('body')}\n\n" + if res.get("source"): + markdown_result += f"*Source: {res.get('source')}*\n\n" + if res.get("image"): + markdown_result += f"![{res.get('title', '')}]({res.get('image')})\n\n" + markdown_result += f"[Read more]({res.get('url', '')})\n\n---\n\n" + + json_result.append( + self.create_json_message( + { + "title": res.get("title", ""), + "date": res.get("date", ""), + "body": res.get("body", ""), + "url": res.get("url", ""), + "image": res.get("image", ""), + "source": res.get("source", ""), + } + ) + ) + + return [self.create_text_message(markdown_result)] + json_result + + def summary_results(self, user_id: str, content: str, query: str) -> str: + prompt = SUMMARY_PROMPT.format(query=query, content=content) + summary = self.invoke_model( + user_id=user_id, + prompt_messages=[ + SystemPromptMessage(content=prompt), + ], + stop=[], + ) + return summary.message.content diff --git a/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_news.yaml b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_news.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6e181e0f41c22fa9128926e7d022eca0e5a2dfa1 --- /dev/null +++ b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_news.yaml @@ -0,0 +1,82 @@ +identity: + name: ddgo_news + author: Assistant + label: + en_US: DuckDuckGo News Search + zh_Hans: DuckDuckGo 新闻搜索 +description: + human: + en_US: Perform news searches on DuckDuckGo and get results. + zh_Hans: 在 DuckDuckGo 上进行新闻搜索并获取结果。 + llm: Perform news searches on DuckDuckGo and get results. +parameters: + - name: query + type: string + required: true + label: + en_US: Query String + zh_Hans: 查询语句 + human_description: + en_US: Search Query. + zh_Hans: 搜索查询语句。 + llm_description: Key words for searching + form: llm + - name: max_results + type: number + required: true + default: 5 + label: + en_US: Max Results + zh_Hans: 最大结果数量 + human_description: + en_US: The Max Results + zh_Hans: 最大结果数量 + form: form + - name: timelimit + type: select + required: false + options: + - value: Day + label: + en_US: Current Day + zh_Hans: 当天 + - value: Week + label: + en_US: Current Week + zh_Hans: 本周 + - value: Month + label: + en_US: Current Month + zh_Hans: 当月 + - value: Year + label: + en_US: Current Year + zh_Hans: 今年 + label: + en_US: Result Time Limit + zh_Hans: 结果时间限制 + human_description: + en_US: Use when querying results within a specific time range only. + zh_Hans: 只查询一定时间范围内的结果时使用 + form: form + - name: require_summary + type: boolean + default: false + label: + en_US: Require Summary + zh_Hans: 是否总结 + human_description: + en_US: Whether to pass the news results to llm for summarization. + zh_Hans: 是否需要将新闻结果传给大模型总结 + form: form + - name: query_prefix + label: + en_US: Query Prefix + zh_Hans: 查询前缀 + type: string + required: false + default: "" + form: form + human_description: + en_US: Specific Search e.g. "site:msn.com" + zh_Hans: 定向搜索 e.g. "site:msn.com" diff --git a/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_search.py b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_search.py new file mode 100644 index 0000000000000000000000000000000000000000..3cd35d16a6f4609e55c29dafba2a219405f70cac --- /dev/null +++ b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_search.py @@ -0,0 +1,50 @@ +from typing import Any + +from duckduckgo_search import DDGS + +from core.model_runtime.entities.message_entities import SystemPromptMessage +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +SUMMARY_PROMPT = """ +User's query: +{query} + +Here is the search engine result: +{content} + +Please summarize the result in a few sentences. +""" + + +class DuckDuckGoSearchTool(BuiltinTool): + """ + Tool for performing a search using DuckDuckGo search engine. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage | list[ToolInvokeMessage]: + query = tool_parameters.get("query") + max_results = tool_parameters.get("max_results", 5) + require_summary = tool_parameters.get("require_summary", False) + + # Add query_prefix handling + query_prefix = tool_parameters.get("query_prefix", "").strip() + final_query = f"{query_prefix} {query}".strip() + + response = DDGS().text(final_query, max_results=max_results) + if require_summary: + results = "\n".join([res.get("body") for res in response]) + results = self.summary_results(user_id=user_id, content=results, query=query) + return self.create_text_message(text=results) + return [self.create_json_message(res) for res in response] + + def summary_results(self, user_id: str, content: str, query: str) -> str: + prompt = SUMMARY_PROMPT.format(query=query, content=content) + summary = self.invoke_model( + user_id=user_id, + prompt_messages=[ + SystemPromptMessage(content=prompt), + ], + stop=[], + ) + return summary.message.content diff --git a/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_search.yaml b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_search.yaml new file mode 100644 index 0000000000000000000000000000000000000000..54e27d9905da1236dd3e6486fcf68cb67294796b --- /dev/null +++ b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_search.yaml @@ -0,0 +1,52 @@ +identity: + name: ddgo_search + author: Yash Parmar + label: + en_US: DuckDuckGo Search + zh_Hans: DuckDuckGo 搜索 +description: + human: + en_US: Perform searches on DuckDuckGo and get results. + zh_Hans: 在 DuckDuckGo 上进行搜索并获取结果。 + llm: Perform searches on DuckDuckGo and get results. +parameters: + - name: query + type: string + required: true + label: + en_US: Query string + zh_Hans: 查询语句 + human_description: + en_US: The search query. + zh_Hans: 搜索查询语句。 + llm_description: Key words for searching + form: llm + - name: max_results + type: number + required: true + default: 5 + label: + en_US: Max results + zh_Hans: 最大结果数量 + form: form + - name: require_summary + type: boolean + default: false + label: + en_US: Require Summary + zh_Hans: 是否总结 + human_description: + en_US: Whether to pass the search results to llm for summarization. + zh_Hans: 是否需要将搜索结果传给大模型总结 + form: form + - name: query_prefix + label: + en_US: Query Prefix + zh_Hans: 查询前缀 + type: string + required: false + default: "" + form: form + human_description: + en_US: Specific Search e.g. "site:wikipedia.org" + zh_Hans: 定向搜索 e.g. "site:wikipedia.org" diff --git a/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_translate.py b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_translate.py new file mode 100644 index 0000000000000000000000000000000000000000..396ce21b183afcc6fd6c7139d01fcdafc0c5c795 --- /dev/null +++ b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_translate.py @@ -0,0 +1,20 @@ +from typing import Any + +from duckduckgo_search import DDGS + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class DuckDuckGoTranslateTool(BuiltinTool): + """ + Tool for performing a search using DuckDuckGo search engine. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + query_dict = { + "keywords": tool_parameters.get("query"), + "to": tool_parameters.get("translate_to"), + } + response = DDGS().translate(**query_dict)[0].get("translated", "Unable to translate!") + return self.create_text_message(text=response) diff --git a/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_translate.yaml b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_translate.yaml new file mode 100644 index 0000000000000000000000000000000000000000..78b5d0b02275b2be13fbc23f00748cd2333dbdd1 --- /dev/null +++ b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_translate.yaml @@ -0,0 +1,51 @@ +identity: + name: ddgo_translate + author: hjlarry + label: + en_US: DuckDuckGo Translate + zh_Hans: DuckDuckGo 翻译 +description: + human: + en_US: Use DuckDuckGo's translation feature. + zh_Hans: 使用DuckDuckGo的翻译功能。 + llm: Use DuckDuckGo's translation feature. +parameters: + - name: query + type: string + required: true + label: + en_US: Translate Content + zh_Hans: 翻译内容 + human_description: + en_US: The translate content. + zh_Hans: 要翻译的内容。 + llm_description: Key words for translate + form: llm + - name: translate_to + type: select + required: true + options: + - value: en + label: + en_US: English + zh_Hans: 英语 + - value: zh-Hans + label: + en_US: Simplified Chinese + zh_Hans: 简体中文 + - value: zh-Hant + label: + en_US: Traditional Chinese + zh_Hans: 繁体中文 + - value: ja + label: + en_US: Japanese + zh_Hans: 日语 + default: en + label: + en_US: Choose Language + zh_Hans: 选择语言 + human_description: + en_US: select the language to translate. + zh_Hans: 选择要翻译的语言 + form: form diff --git a/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_video.py b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_video.py new file mode 100644 index 0000000000000000000000000000000000000000..1eef0b1ba23d42ae584729ec1a9725ef391fbf06 --- /dev/null +++ b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_video.py @@ -0,0 +1,91 @@ +from typing import Any, ClassVar + +from duckduckgo_search import DDGS + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class DuckDuckGoVideoSearchTool(BuiltinTool): + """ + Tool for performing a video search using DuckDuckGo search engine. + """ + + IFRAME_TEMPLATE: ClassVar[str] = """ +
+ +
""" + + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> list[ToolInvokeMessage]: + query_dict = { + "keywords": tool_parameters.get("query"), # LLM's query + "region": tool_parameters.get("region", "wt-wt"), + "safesearch": tool_parameters.get("safesearch", "moderate"), + "timelimit": tool_parameters.get("timelimit"), + "resolution": tool_parameters.get("resolution"), + "duration": tool_parameters.get("duration"), + "license_videos": tool_parameters.get("license_videos"), + "max_results": tool_parameters.get("max_results"), + } + + # Remove None values to use API defaults + query_dict = {k: v for k, v in query_dict.items() if v is not None} + + # Get proxy URL from parameters + proxy_url = tool_parameters.get("proxy_url", "").strip() + + query_prefix = tool_parameters.get("query_prefix", "").strip() + final_query = f"{query_prefix} {query_dict['keywords']}".strip() + + # Update the keywords in query_dict with the final_query + query_dict["keywords"] = final_query + + response = DDGS().videos(**query_dict) + + # Create HTML result with embedded iframes + markdown_result = "\n\n" + json_result = [] + + for res in response: + title = res.get("title", "") + embed_html = res.get("embed_html", "") + description = res.get("description", "") + content_url = res.get("content", "") + transcript_url = None + + # Handle TED.com videos + if "ted.com/talks" in content_url: + # Create transcript URL + transcript_url = f"{content_url}/transcript" + # Create embed URL + embed_url = content_url.replace("www.ted.com", "embed.ted.com") + if proxy_url: + embed_url = f"{proxy_url}{embed_url}" + embed_html = self.IFRAME_TEMPLATE.format(src=embed_url) + + # Original YouTube/other platform handling + elif embed_html: + embed_url = res.get("embed_url", "") + if proxy_url and embed_url: + embed_url = f"{proxy_url}{embed_url}" + embed_html = self.IFRAME_TEMPLATE.format(src=embed_url) + + markdown_result += f"{title}\n\n" + markdown_result += f"{embed_html}\n\n" + if description: + markdown_result += f"{description}\n\n" + markdown_result += "---\n\n" + + # Add transcript_url to the JSON result if available + result_dict = res.copy() + if transcript_url: + result_dict["transcript_url"] = transcript_url + json_result.append(self.create_json_message(result_dict)) + + return [self.create_text_message(markdown_result)] + json_result diff --git a/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_video.yaml b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_video.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d846244e3dfcbd14b941fe2189bb92cfa0bc47c2 --- /dev/null +++ b/api/core/tools/provider/builtin/duckduckgo/tools/ddgo_video.yaml @@ -0,0 +1,108 @@ +identity: + name: ddgo_video + author: Tao Wang + label: + en_US: DuckDuckGo Video Search + zh_Hans: DuckDuckGo 视频搜索 +description: + human: + en_US: Search and embedded videos. + zh_Hans: 搜索并嵌入视频 + llm: Search videos on duckduckgo and embed videos in iframe +parameters: + - name: query + label: + en_US: Query String + zh_Hans: 查询语句 + type: string + required: true + human_description: + en_US: Search Query + zh_Hans: 搜索查询语句 + llm_description: Key words for searching + form: llm + - name: max_results + label: + en_US: Max Results + zh_Hans: 最大结果数量 + type: number + required: true + default: 3 + minimum: 1 + maximum: 10 + human_description: + en_US: The max results (1-10) + zh_Hans: 最大结果数量(1-10) + form: form + - name: timelimit + label: + en_US: Result Time Limit + zh_Hans: 结果时间限制 + type: select + required: false + options: + - value: Day + label: + en_US: Current Day + zh_Hans: 当天 + - value: Week + label: + en_US: Current Week + zh_Hans: 本周 + - value: Month + label: + en_US: Current Month + zh_Hans: 当月 + - value: Year + label: + en_US: Current Year + zh_Hans: 今年 + human_description: + en_US: Query results within a specific time range only + zh_Hans: 只查询一定时间范围内的结果时使用 + form: form + - name: duration + label: + en_US: Video Duration + zh_Hans: 视频时长 + type: select + required: false + options: + - value: short + label: + en_US: Short (<4 minutes) + zh_Hans: 短视频(<4分钟) + - value: medium + label: + en_US: Medium (4-20 minutes) + zh_Hans: 中等(4-20分钟) + - value: long + label: + en_US: Long (>20 minutes) + zh_Hans: 长视频(>20分钟) + human_description: + en_US: Filter videos by duration + zh_Hans: 按时长筛选视频 + form: form + - name: proxy_url + label: + en_US: Proxy URL + zh_Hans: 视频代理地址 + type: string + required: false + default: "" + human_description: + en_US: Proxy URL + zh_Hans: 视频代理地址 + form: form + - name: query_prefix + label: + en_US: Query Prefix + zh_Hans: 查询前缀 + type: string + required: false + default: "" + form: form + human_description: + en_US: Specific Search e.g. "site:www.ted.com" + zh_Hans: 定向搜索 e.g. "site:www.ted.com" diff --git a/api/core/tools/provider/builtin/email/_assets/icon.svg b/api/core/tools/provider/builtin/email/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..b34f333890bc1aea2b299e38f7b747528ba14d0d --- /dev/null +++ b/api/core/tools/provider/builtin/email/_assets/icon.svg @@ -0,0 +1 @@ + diff --git a/api/core/tools/provider/builtin/email/email.py b/api/core/tools/provider/builtin/email/email.py new file mode 100644 index 0000000000000000000000000000000000000000..182d8dac28efea1954555cb0fa773ff00ec1b2a8 --- /dev/null +++ b/api/core/tools/provider/builtin/email/email.py @@ -0,0 +1,7 @@ +from core.tools.provider.builtin.email.tools.send_mail import SendMailTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class SmtpProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + SendMailTool() diff --git a/api/core/tools/provider/builtin/email/email.yaml b/api/core/tools/provider/builtin/email/email.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bb1bb7f6f3e972c3dec12208d89ae7c964866740 --- /dev/null +++ b/api/core/tools/provider/builtin/email/email.yaml @@ -0,0 +1,83 @@ +identity: + author: wakaka6 + name: email + label: + en_US: email + zh_Hans: 电子邮件 + description: + en_US: send email through smtp protocol + zh_Hans: 通过smtp协议发送电子邮件 + icon: icon.svg + tags: + - utilities +credentials_for_provider: + email_account: + type: text-input + required: true + label: + en_US: email account + zh_Hans: 邮件账号 + placeholder: + en_US: input you email account + zh_Hans: 输入你的邮箱账号 + help: + en_US: email account + zh_Hans: 邮件账号 + email_password: + type: secret-input + required: true + label: + en_US: email password + zh_Hans: 邮件密码 + placeholder: + en_US: email password + zh_Hans: 邮件密码 + help: + en_US: email password + zh_Hans: 邮件密码 + smtp_server: + type: text-input + required: true + label: + en_US: smtp server + zh_Hans: 发信smtp服务器地址 + placeholder: + en_US: smtp server + zh_Hans: 发信smtp服务器地址 + help: + en_US: smtp server + zh_Hans: 发信smtp服务器地址 + smtp_port: + type: text-input + required: true + label: + en_US: smtp server port + zh_Hans: 发信smtp服务器端口 + placeholder: + en_US: smtp server port + zh_Hans: 发信smtp服务器端口 + help: + en_US: smtp server port + zh_Hans: 发信smtp服务器端口 + encrypt_method: + type: select + required: true + options: + - value: NONE + label: + en_US: NONE + zh_Hans: 无加密 + - value: SSL + label: + en_US: SSL + zh_Hans: SSL加密 + - value: TLS + label: + en_US: START TLS + zh_Hans: START TLS加密 + label: + en_US: encrypt method + zh_Hans: 加密方式 + help: + en_US: smtp server encrypt method + zh_Hans: 发信smtp服务器加密方式 diff --git a/api/core/tools/provider/builtin/email/tools/send.py b/api/core/tools/provider/builtin/email/tools/send.py new file mode 100644 index 0000000000000000000000000000000000000000..2012d8b1156fb52792e14ba25c948a9a714c7411 --- /dev/null +++ b/api/core/tools/provider/builtin/email/tools/send.py @@ -0,0 +1,53 @@ +import logging +import smtplib +import ssl +from email.mime.multipart import MIMEMultipart +from email.mime.text import MIMEText + +from pydantic import BaseModel + + +class SendEmailToolParameters(BaseModel): + smtp_server: str + smtp_port: int + + email_account: str + email_password: str + + sender_to: str + subject: str + email_content: str + encrypt_method: str + + +def send_mail(params: SendEmailToolParameters): + timeout = 60 + msg = MIMEMultipart("alternative") + msg["From"] = params.email_account + msg["To"] = params.sender_to + msg["Subject"] = params.subject + msg.attach(MIMEText(params.email_content, "plain")) + msg.attach(MIMEText(params.email_content, "html")) + + ctx = ssl.create_default_context() + + if params.encrypt_method.upper() == "SSL": + try: + with smtplib.SMTP_SSL(params.smtp_server, params.smtp_port, context=ctx, timeout=timeout) as server: + server.login(params.email_account, params.email_password) + server.sendmail(params.email_account, params.sender_to, msg.as_string()) + return True + except Exception as e: + logging.exception("send email failed") + return False + else: # NONE or TLS + try: + with smtplib.SMTP(params.smtp_server, params.smtp_port, timeout=timeout) as server: + if params.encrypt_method.upper() == "TLS": + server.starttls(context=ctx) + server.login(params.email_account, params.email_password) + server.sendmail(params.email_account, params.sender_to, msg.as_string()) + return True + except Exception as e: + logging.exception("send email failed") + return False diff --git a/api/core/tools/provider/builtin/email/tools/send_mail.py b/api/core/tools/provider/builtin/email/tools/send_mail.py new file mode 100644 index 0000000000000000000000000000000000000000..33c040400ca14a98ffe510b2caab2dbf298f1d88 --- /dev/null +++ b/api/core/tools/provider/builtin/email/tools/send_mail.py @@ -0,0 +1,66 @@ +import re +from typing import Any, Union + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.email.tools.send import ( + SendEmailToolParameters, + send_mail, +) +from core.tools.tool.builtin_tool import BuiltinTool + + +class SendMailTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + sender = self.runtime.credentials.get("email_account", "") + email_rgx = re.compile(r"^[a-zA-Z0-9._-]+@[a-zA-Z0-9_-]+(\.[a-zA-Z0-9_-]+)+$") + password = self.runtime.credentials.get("email_password", "") + smtp_server = self.runtime.credentials.get("smtp_server", "") + if not smtp_server: + return self.create_text_message("please input smtp server") + smtp_port = self.runtime.credentials.get("smtp_port", "") + try: + smtp_port = int(smtp_port) + except ValueError: + return self.create_text_message("Invalid parameter smtp_port(should be int)") + + if not sender: + return self.create_text_message("please input sender") + if not email_rgx.match(sender): + return self.create_text_message("Invalid parameter userid, the sender is not a mailbox") + + receiver_email = tool_parameters["send_to"] + if not receiver_email: + return self.create_text_message("please input receiver email") + if not email_rgx.match(receiver_email): + return self.create_text_message("Invalid parameter receiver email, the receiver email is not a mailbox") + email_content = tool_parameters.get("email_content", "") + + if not email_content: + return self.create_text_message("please input email content") + + subject = tool_parameters.get("subject", "") + if not subject: + return self.create_text_message("please input email subject") + + encrypt_method = self.runtime.credentials.get("encrypt_method", "") + if not encrypt_method: + return self.create_text_message("please input encrypt method") + + send_email_params = SendEmailToolParameters( + smtp_server=smtp_server, + smtp_port=smtp_port, + email_account=sender, + email_password=password, + sender_to=receiver_email, + subject=subject, + email_content=email_content, + encrypt_method=encrypt_method, + ) + if send_mail(send_email_params): + return self.create_text_message("send email success") + return self.create_text_message("send email failed") diff --git a/api/core/tools/provider/builtin/email/tools/send_mail.yaml b/api/core/tools/provider/builtin/email/tools/send_mail.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f54880bf3e9e22cca1f8c0b302148898922c3c92 --- /dev/null +++ b/api/core/tools/provider/builtin/email/tools/send_mail.yaml @@ -0,0 +1,46 @@ +identity: + name: send_mail + author: wakaka6 + label: + en_US: send email + zh_Hans: 发送邮件 + icon: icon.svg +description: + human: + en_US: A tool for sending email + zh_Hans: 用于发送邮件 + llm: A tool for sending email +parameters: + - name: send_to + type: string + required: true + label: + en_US: Recipient email account + zh_Hans: 收件人邮箱账号 + human_description: + en_US: Recipient email account + zh_Hans: 收件人邮箱账号 + llm_description: Recipient email account + form: llm + - name: subject + type: string + required: true + label: + en_US: email subject + zh_Hans: 邮件主题 + human_description: + en_US: email subject + zh_Hans: 邮件主题 + llm_description: email subject + form: llm + - name: email_content + type: string + required: true + label: + en_US: email content + zh_Hans: 邮件内容 + human_description: + en_US: email content + zh_Hans: 邮件内容 + llm_description: email content + form: llm diff --git a/api/core/tools/provider/builtin/email/tools/send_mail_batch.py b/api/core/tools/provider/builtin/email/tools/send_mail_batch.py new file mode 100644 index 0000000000000000000000000000000000000000..537dedb27d530f25f1151cc67e6b2b0cfb6e18a6 --- /dev/null +++ b/api/core/tools/provider/builtin/email/tools/send_mail_batch.py @@ -0,0 +1,75 @@ +import json +import re +from typing import Any, Union + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.email.tools.send import ( + SendEmailToolParameters, + send_mail, +) +from core.tools.tool.builtin_tool import BuiltinTool + + +class SendMailTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + sender = self.runtime.credentials.get("email_account", "") + email_rgx = re.compile(r"^[a-zA-Z0-9._-]+@[a-zA-Z0-9_-]+(\.[a-zA-Z0-9_-]+)+$") + password = self.runtime.credentials.get("email_password", "") + smtp_server = self.runtime.credentials.get("smtp_server", "") + if not smtp_server: + return self.create_text_message("please input smtp server") + smtp_port = self.runtime.credentials.get("smtp_port", "") + try: + smtp_port = int(smtp_port) + except ValueError: + return self.create_text_message("Invalid parameter smtp_port(should be int)") + + if not sender: + return self.create_text_message("please input sender") + if not email_rgx.match(sender): + return self.create_text_message("Invalid parameter userid, the sender is not a mailbox") + + receivers_email = tool_parameters["send_to"] + if not receivers_email: + return self.create_text_message("please input receiver email") + receivers_email = json.loads(receivers_email) + for receiver in receivers_email: + if not email_rgx.match(receiver): + return self.create_text_message( + f"Invalid parameter receiver email, the receiver email({receiver}) is not a mailbox" + ) + email_content = tool_parameters.get("email_content", "") + + if not email_content: + return self.create_text_message("please input email content") + + subject = tool_parameters.get("subject", "") + if not subject: + return self.create_text_message("please input email subject") + + encrypt_method = self.runtime.credentials.get("encrypt_method", "") + if not encrypt_method: + return self.create_text_message("please input encrypt method") + + msg = {} + for receiver in receivers_email: + send_email_params = SendEmailToolParameters( + smtp_server=smtp_server, + smtp_port=smtp_port, + email_account=sender, + email_password=password, + sender_to=receiver, + subject=subject, + email_content=email_content, + encrypt_method=encrypt_method, + ) + if send_mail(send_email_params): + msg[receiver] = "send email success" + else: + msg[receiver] = "send email failed" + return self.create_text_message(json.dumps(msg)) diff --git a/api/core/tools/provider/builtin/email/tools/send_mail_batch.yaml b/api/core/tools/provider/builtin/email/tools/send_mail_batch.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6e4aa785cb946e1665a81d27934a67dfdde8f044 --- /dev/null +++ b/api/core/tools/provider/builtin/email/tools/send_mail_batch.yaml @@ -0,0 +1,46 @@ +identity: + name: send_mail_batch + author: wakaka6 + label: + en_US: send email to multiple recipients + zh_Hans: 发送邮件给多个收件人 + icon: icon.svg +description: + human: + en_US: A tool for sending email to multiple recipients + zh_Hans: 用于发送邮件给多个收件人的工具 + llm: A tool for sending email to multiple recipients +parameters: + - name: send_to + type: string + required: true + label: + en_US: Recipient email account(json list) + zh_Hans: 收件人邮箱账号(json list) + human_description: + en_US: Recipient email account + zh_Hans: 收件人邮箱账号 + llm_description: A list of recipient email account(json format) + form: llm + - name: subject + type: string + required: true + label: + en_US: email subject + zh_Hans: 邮件主题 + human_description: + en_US: email subject + zh_Hans: 邮件主题 + llm_description: email subject + form: llm + - name: email_content + type: string + required: true + label: + en_US: email content + zh_Hans: 邮件内容 + human_description: + en_US: email content + zh_Hans: 邮件内容 + llm_description: email content + form: llm diff --git a/api/core/tools/provider/builtin/fal/_assets/icon.svg b/api/core/tools/provider/builtin/fal/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..bfb270774dd14c79ed820138c42e61f1394c8fd0 --- /dev/null +++ b/api/core/tools/provider/builtin/fal/_assets/icon.svg @@ -0,0 +1,4 @@ + + + + diff --git a/api/core/tools/provider/builtin/fal/fal.py b/api/core/tools/provider/builtin/fal/fal.py new file mode 100644 index 0000000000000000000000000000000000000000..c68e2021331082d9eaf946f5616cc45fa97c346b --- /dev/null +++ b/api/core/tools/provider/builtin/fal/fal.py @@ -0,0 +1,20 @@ +import requests + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class FalProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + url = "https://fal.run/fal-ai/flux/dev" + headers = { + "Authorization": f"Key {credentials.get('fal_api_key')}", + "Content-Type": "application/json", + } + data = {"prompt": "Cat"} + + response = requests.post(url, json=data, headers=headers) + if response.status_code == 401: + raise ToolProviderCredentialValidationError("FAL API key is invalid") + elif response.status_code != 200: + raise ToolProviderCredentialValidationError(f"FAL API key validation failed: {response.text}") diff --git a/api/core/tools/provider/builtin/fal/fal.yaml b/api/core/tools/provider/builtin/fal/fal.yaml new file mode 100644 index 0000000000000000000000000000000000000000..050a73f62660f6a5d68816675fbf4bd67bc33b6b --- /dev/null +++ b/api/core/tools/provider/builtin/fal/fal.yaml @@ -0,0 +1,21 @@ +identity: + author: Kalo Chin + name: fal + label: + en_US: FAL + zh_CN: FAL + description: + en_US: The image generation API provided by FAL. + zh_CN: FAL 提供的图像生成 API。 + icon: icon.svg + tags: + - image +credentials_for_provider: + fal_api_key: + type: secret-input + required: true + label: + en_US: FAL API Key + placeholder: + en_US: Please input your FAL API key + url: https://fal.ai/dashboard/keys diff --git a/api/core/tools/provider/builtin/fal/tools/flux_1_1_pro.py b/api/core/tools/provider/builtin/fal/tools/flux_1_1_pro.py new file mode 100644 index 0000000000000000000000000000000000000000..7b5f10a64d4cdd8ef19a1c0355a21c338a375916 --- /dev/null +++ b/api/core/tools/provider/builtin/fal/tools/flux_1_1_pro.py @@ -0,0 +1,46 @@ +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class Flux11ProTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + headers = { + "Authorization": f"Key {self.runtime.credentials['fal_api_key']}", + "Content-Type": "application/json", + } + + prompt = tool_parameters.get("prompt", "") + sanitized_prompt = prompt.replace("\\", "") # Remove backslashes from the prompt which may cause errors + + payload = { + "prompt": sanitized_prompt, + "image_size": tool_parameters.get("image_size", "landscape_4_3"), + "seed": tool_parameters.get("seed"), + "sync_mode": tool_parameters.get("sync_mode", False), + "num_images": tool_parameters.get("num_images", 1), + "enable_safety_checker": tool_parameters.get("enable_safety_checker", True), + "safety_tolerance": tool_parameters.get("safety_tolerance", "2"), + } + + url = "https://fal.run/fal-ai/flux-pro/v1.1" + + response = requests.post(url, json=payload, headers=headers) + + if response.status_code != 200: + return self.create_text_message(f"Got Error Response: {response.text}") + + res = response.json() + result = [self.create_json_message(res)] + + for image_info in res.get("images", []): + image_url = image_info.get("url") + if image_url: + result.append(self.create_image_message(image=image_url, save_as=self.VariableKey.IMAGE.value)) + + return result diff --git a/api/core/tools/provider/builtin/fal/tools/flux_1_1_pro.yaml b/api/core/tools/provider/builtin/fal/tools/flux_1_1_pro.yaml new file mode 100644 index 0000000000000000000000000000000000000000..237ee9937f925bb9fbd79428b856fc1241ba4e5f --- /dev/null +++ b/api/core/tools/provider/builtin/fal/tools/flux_1_1_pro.yaml @@ -0,0 +1,147 @@ +identity: + name: flux_1_1_pro + author: Kalo Chin + label: + en_US: FLUX 1.1 [pro] + zh_Hans: FLUX 1.1 [pro] + icon: icon.svg +description: + human: + en_US: FLUX 1.1 [pro] is an enhanced version of FLUX.1 [pro], improved image generation capabilities, delivering superior composition, detail, and artistic fidelity compared to its predecessor. + zh_Hans: FLUX 1.1 [pro] 是 FLUX.1 [pro] 的增强版,改进了图像生成能力,与其前身相比,提供了更出色的构图、细节和艺术保真度。 + llm: This tool generates images from prompts using FAL's FLUX 1.1 [pro] model. +parameters: + - name: prompt + type: string + required: true + label: + en_US: Prompt + zh_Hans: 提示词 + human_description: + en_US: The text prompt used to generate the image. + zh_Hans: 用于生成图片的文字提示词。 + llm_description: This prompt text will be used to generate the image. + form: llm + - name: image_size + type: select + required: false + options: + - value: square_hd + label: + en_US: Square HD + zh_Hans: 方形高清 + - value: square + label: + en_US: Square + zh_Hans: 方形 + - value: portrait_4_3 + label: + en_US: Portrait 4:3 + zh_Hans: 竖屏 4:3 + - value: portrait_16_9 + label: + en_US: Portrait 16:9 + zh_Hans: 竖屏 16:9 + - value: landscape_4_3 + label: + en_US: Landscape 4:3 + zh_Hans: 横屏 4:3 + - value: landscape_16_9 + label: + en_US: Landscape 16:9 + zh_Hans: 横屏 16:9 + default: landscape_4_3 + label: + en_US: Image Size + zh_Hans: 图片大小 + human_description: + en_US: The size of the generated image. + zh_Hans: 生成图像的尺寸。 + form: form + - name: num_images + type: number + required: false + default: 1 + min: 1 + max: 1 + label: + en_US: Number of Images + zh_Hans: 图片数量 + human_description: + en_US: The number of images to generate. + zh_Hans: 要生成的图片数量。 + form: form + - name: safety_tolerance + type: select + required: false + options: + - value: "1" + label: + en_US: "1 (Most strict)" + zh_Hans: "1(最严格)" + - value: "2" + label: + en_US: "2" + zh_Hans: "2" + - value: "3" + label: + en_US: "3" + zh_Hans: "3" + - value: "4" + label: + en_US: "4" + zh_Hans: "4" + - value: "5" + label: + en_US: "5" + zh_Hans: "5" + - value: "6" + label: + en_US: "6 (Most permissive)" + zh_Hans: "6(最宽松)" + default: "2" + label: + en_US: Safety Tolerance + zh_Hans: 安全容忍度 + human_description: + en_US: The safety tolerance level for the generated image. 1 being the most strict and 6 being the most permissive. + zh_Hans: 生成图像的安全容忍级别,1 为最严格,6 为最宽松。 + form: form + - name: seed + type: number + required: false + min: 0 + max: 9999999999 + label: + en_US: Seed + zh_Hans: 种子 + human_description: + en_US: The same seed and prompt can produce similar images. + zh_Hans: 相同的种子和提示词可以产生相似的图像。 + form: form + - name: enable_safety_checker + type: boolean + required: false + default: true + label: + en_US: Enable Safety Checker + zh_Hans: 启用安全检查器 + human_description: + en_US: Enable or disable the safety checker. + zh_Hans: 启用或禁用安全检查器。 + form: form + - name: sync_mode + type: boolean + required: false + default: false + label: + en_US: Sync Mode + zh_Hans: 同步模式 + human_description: + en_US: > + If set to true, the function will wait for the image to be generated and uploaded before returning the response. + This will increase the latency but allows you to get the image directly in the response without going through the CDN. + zh_Hans: > + 如果设置为 true,函数将在生成并上传图像后再返回响应。 + 这将增加函数的延迟,但可以让您直接在响应中获取图像,而无需通过 CDN。 + form: form diff --git a/api/core/tools/provider/builtin/fal/tools/flux_1_1_pro_ultra.py b/api/core/tools/provider/builtin/fal/tools/flux_1_1_pro_ultra.py new file mode 100644 index 0000000000000000000000000000000000000000..2fb1565e7cd1b2fd16f48467b2fe56b6c68b827a --- /dev/null +++ b/api/core/tools/provider/builtin/fal/tools/flux_1_1_pro_ultra.py @@ -0,0 +1,47 @@ +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class Flux11ProUltraTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + headers = { + "Authorization": f"Key {self.runtime.credentials['fal_api_key']}", + "Content-Type": "application/json", + } + + prompt = tool_parameters.get("prompt", "") + sanitized_prompt = prompt.replace("\\", "") # Remove backslashes from the prompt which may cause errors + + payload = { + "prompt": sanitized_prompt, + "seed": tool_parameters.get("seed"), + "sync_mode": tool_parameters.get("sync_mode", False), + "num_images": tool_parameters.get("num_images", 1), + "enable_safety_checker": tool_parameters.get("enable_safety_checker", True), + "safety_tolerance": str(tool_parameters.get("safety_tolerance", "2")), + "aspect_ratio": tool_parameters.get("aspect_ratio", "16:9"), + "raw": tool_parameters.get("raw", False), + } + + url = "https://fal.run/fal-ai/flux-pro/v1.1-ultra" + + response = requests.post(url, json=payload, headers=headers) + + if response.status_code != 200: + return self.create_text_message(f"Got Error Response: {response.text}") + + res = response.json() + result = [self.create_json_message(res)] + + for image_info in res.get("images", []): + image_url = image_info.get("url") + if image_url: + result.append(self.create_image_message(image=image_url, save_as=self.VariableKey.IMAGE.value)) + + return result diff --git a/api/core/tools/provider/builtin/fal/tools/flux_1_1_pro_ultra.yaml b/api/core/tools/provider/builtin/fal/tools/flux_1_1_pro_ultra.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d518e5192935b96087d28f3eb4769274481e2571 --- /dev/null +++ b/api/core/tools/provider/builtin/fal/tools/flux_1_1_pro_ultra.yaml @@ -0,0 +1,162 @@ +identity: + name: flux_1_1_pro_ultra + author: Kalo Chin + label: + en_US: FLUX 1.1 [pro] ultra + zh_Hans: FLUX 1.1 [pro] ultra + icon: icon.svg +description: + human: + en_US: FLUX 1.1 [pro] ultra is the newest version of FLUX 1.1 [pro], maintaining professional-grade image quality while delivering up to 2K resolution with improved photo realism. + zh_Hans: FLUX 1.1 [pro] ultra 是 FLUX 1.1 [pro] 的最新版本,保持了专业级的图像质量,同时以改进的照片真实感提供高达 2K 的分辨率。 + llm: This tool generates images from prompts using FAL's FLUX 1.1 [pro] ultra model. +parameters: + - name: prompt + type: string + required: true + label: + en_US: Prompt + zh_Hans: 提示词 + human_description: + en_US: The text prompt used to generate the image. + zh_Hans: 用于生成图像的文本提示。 + llm_description: This prompt text will be used to generate the image. + form: llm + - name: aspect_ratio + type: select + required: false + options: + - value: '21:9' + label: + en_US: '21:9' + zh_Hans: '21:9' + - value: '16:9' + label: + en_US: '16:9' + zh_Hans: '16:9' + - value: '4:3' + label: + en_US: '4:3' + zh_Hans: '4:3' + - value: '1:1' + label: + en_US: '1:1' + zh_Hans: '1:1' + - value: '3:4' + label: + en_US: '3:4' + zh_Hans: '3:4' + - value: '9:16' + label: + en_US: '9:16' + zh_Hans: '9:16' + - value: '9:21' + label: + en_US: '9:21' + zh_Hans: '9:21' + default: '16:9' + label: + en_US: Aspect Ratio + zh_Hans: 纵横比 + human_description: + en_US: The aspect ratio of the generated image. + zh_Hans: 生成图像的宽高比。 + form: form + - name: num_images + type: number + required: false + default: 1 + min: 1 + max: 1 + label: + en_US: Number of Images + zh_Hans: 图片数量 + human_description: + en_US: The number of images to generate. + zh_Hans: 要生成的图像数量。 + form: form + - name: safety_tolerance + type: select + required: false + options: + - value: "1" + label: + en_US: "1 (Most strict)" + zh_Hans: "1(最严格)" + - value: "2" + label: + en_US: "2" + zh_Hans: "2" + - value: "3" + label: + en_US: "3" + zh_Hans: "3" + - value: "4" + label: + en_US: "4" + zh_Hans: "4" + - value: "5" + label: + en_US: "5" + zh_Hans: "5" + - value: "6" + label: + en_US: "6 (Most permissive)" + zh_Hans: "6(最宽松)" + default: '2' + label: + en_US: Safety Tolerance + zh_Hans: 安全容忍度 + human_description: + en_US: The safety tolerance level for the generated image. 1 being the most strict and 6 being the most permissive. + zh_Hans: 生成图像的安全容忍级别,1 为最严格,6 为最宽松。 + form: form + - name: seed + type: number + required: false + min: 0 + max: 9999999999 + label: + en_US: Seed + zh_Hans: 种子 + human_description: + en_US: The same seed and prompt can produce similar images. + zh_Hans: 相同的种子和提示词可以生成相似的图像。 + form: form + - name: raw + type: boolean + required: false + default: false + label: + en_US: Raw Mode + zh_Hans: 原始模式 + human_description: + en_US: Generate less processed, more natural-looking images. + zh_Hans: 生成较少处理、更自然的图像。 + form: form + - name: enable_safety_checker + type: boolean + required: false + default: true + label: + en_US: Enable Safety Checker + zh_Hans: 启用安全检查器 + human_description: + en_US: Enable or disable the safety checker. + zh_Hans: 启用或禁用安全检查器。 + form: form + - name: sync_mode + type: boolean + required: false + default: false + label: + en_US: Sync Mode + zh_Hans: 同步模式 + human_description: + en_US: > + If set to true, the function will wait for the image to be generated and uploaded before returning the response. + This will increase the latency but allows you to get the image directly in the response without going through the CDN. + zh_Hans: > + 如果设置为 true,函数将在生成并上传图像后才返回响应。 + 这将增加延迟,但允许您直接在响应中获取图像,而无需通过 CDN。 + form: form diff --git a/api/core/tools/provider/builtin/fal/tools/flux_1_dev.py b/api/core/tools/provider/builtin/fal/tools/flux_1_dev.py new file mode 100644 index 0000000000000000000000000000000000000000..b44d9fe752ed5ad847bc4b065dd0dff26f63e04b --- /dev/null +++ b/api/core/tools/provider/builtin/fal/tools/flux_1_dev.py @@ -0,0 +1,47 @@ +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class Flux1DevTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + headers = { + "Authorization": f"Key {self.runtime.credentials['fal_api_key']}", + "Content-Type": "application/json", + } + + prompt = tool_parameters.get("prompt", "") + sanitized_prompt = prompt.replace("\\", "") # Remove backslashes from the prompt which may cause errors + + payload = { + "prompt": sanitized_prompt, + "image_size": tool_parameters.get("image_size", "landscape_4_3"), + "num_inference_steps": tool_parameters.get("num_inference_steps", 28), + "guidance_scale": tool_parameters.get("guidance_scale", 3.5), + "seed": tool_parameters.get("seed"), + "num_images": tool_parameters.get("num_images", 1), + "enable_safety_checker": tool_parameters.get("enable_safety_checker", True), + "sync_mode": tool_parameters.get("sync_mode", False), + } + + url = "https://fal.run/fal-ai/flux/dev" + + response = requests.post(url, json=payload, headers=headers) + + if response.status_code != 200: + return self.create_text_message(f"Got Error Response: {response.text}") + + res = response.json() + result = [self.create_json_message(res)] + + for image_info in res.get("images", []): + image_url = image_info.get("url") + if image_url: + result.append(self.create_image_message(image=image_url, save_as=self.VariableKey.IMAGE.value)) + + return result diff --git a/api/core/tools/provider/builtin/fal/tools/flux_1_dev.yaml b/api/core/tools/provider/builtin/fal/tools/flux_1_dev.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3b22af941fc60ac1455b0bb5ebb0fa14fe01a8b6 --- /dev/null +++ b/api/core/tools/provider/builtin/fal/tools/flux_1_dev.yaml @@ -0,0 +1,137 @@ +identity: + name: flux_1_dev + author: Kalo Chin + label: + en_US: FLUX.1 [dev] + zh_Hans: FLUX.1 [dev] + icon: icon.svg +description: + human: + en_US: FLUX.1 [dev] is a 12 billion parameter flow transformer that generates high-quality images from text. It is suitable for personal and commercial use. + zh_Hans: FLUX.1 [dev] 是一个拥有120亿参数的流动变换模型,可以从文本生成高质量的图像。适用于个人和商业用途。 + llm: This tool generates images from prompts using FAL's FLUX.1 [dev] model. +parameters: + - name: prompt + type: string + required: true + label: + en_US: Prompt + zh_Hans: 提示词 + human_description: + en_US: The text prompt used to generate the image. + zh_Hans: 用于生成图片的文字提示词。 + llm_description: This prompt text will be used to generate the image. + form: llm + - name: image_size + type: select + required: false + options: + - value: square_hd + label: + en_US: Square HD + zh_Hans: 方形高清 + - value: square + label: + en_US: Square + zh_Hans: 方形 + - value: portrait_4_3 + label: + en_US: Portrait 4:3 + zh_Hans: 竖屏 4:3 + - value: portrait_16_9 + label: + en_US: Portrait 16:9 + zh_Hans: 竖屏 16:9 + - value: landscape_4_3 + label: + en_US: Landscape 4:3 + zh_Hans: 横屏 4:3 + - value: landscape_16_9 + label: + en_US: Landscape 16:9 + zh_Hans: 横屏 16:9 + default: landscape_4_3 + label: + en_US: Image Size + zh_Hans: 图片大小 + human_description: + en_US: The size of the generated image. + zh_Hans: 生成图像的尺寸。 + form: form + - name: num_images + type: number + required: false + default: 1 + min: 1 + max: 4 + label: + en_US: Number of Images + zh_Hans: 图片数量 + human_description: + en_US: The number of images to generate. + zh_Hans: 要生成的图片数量。 + form: form + - name: num_inference_steps + type: number + required: false + default: 28 + min: 1 + max: 50 + label: + en_US: Num Inference Steps + zh_Hans: 推理步数 + human_description: + en_US: The number of inference steps to perform. More steps produce higher quality but take longer. + zh_Hans: 执行的推理步骤数量。更多的步骤可以产生更高质量的结果,但需要更长的时间。 + form: form + - name: guidance_scale + type: number + required: false + default: 3.5 + min: 0 + max: 20 + label: + en_US: Guidance Scale + zh_Hans: 指导强度 + human_description: + en_US: How closely the model should follow the prompt. + zh_Hans: 模型对提示词的遵循程度。 + form: form + - name: seed + type: number + required: false + min: 0 + max: 9999999999 + label: + en_US: Seed + zh_Hans: 种子 + human_description: + en_US: The same seed and prompt can produce similar images. + zh_Hans: 相同的种子和提示可以产生相似的图像。 + form: form + - name: enable_safety_checker + type: boolean + required: false + default: true + label: + en_US: Enable Safety Checker + zh_Hans: 启用安全检查器 + human_description: + en_US: Enable or disable the safety checker. + zh_Hans: 启用或禁用安全检查器。 + form: form + - name: sync_mode + type: boolean + required: false + default: false + label: + en_US: Sync Mode + zh_Hans: 同步模式 + human_description: + en_US: > + If set to true, the function will wait for the image to be generated and uploaded before returning the response. + This will increase the latency but allows you to get the image directly in the response without going through the CDN. + zh_Hans: > + 如果设置为 true,函数将在生成并上传图像后再返回响应。 + 这将增加函数的延迟,但可以让您直接在响应中获取图像,而无需通过 CDN。 + form: form diff --git a/api/core/tools/provider/builtin/fal/tools/flux_1_pro_new.py b/api/core/tools/provider/builtin/fal/tools/flux_1_pro_new.py new file mode 100644 index 0000000000000000000000000000000000000000..be60366155dbe32059b6084dd3f14c65a99820e7 --- /dev/null +++ b/api/core/tools/provider/builtin/fal/tools/flux_1_pro_new.py @@ -0,0 +1,47 @@ +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class Flux1ProNewTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + headers = { + "Authorization": f"Key {self.runtime.credentials['fal_api_key']}", + "Content-Type": "application/json", + } + + prompt = tool_parameters.get("prompt", "") + sanitized_prompt = prompt.replace("\\", "") # Remove backslashes that may cause errors + + payload = { + "prompt": sanitized_prompt, + "image_size": tool_parameters.get("image_size", "landscape_4_3"), + "num_inference_steps": tool_parameters.get("num_inference_steps", 28), + "guidance_scale": tool_parameters.get("guidance_scale", 3.5), + "seed": tool_parameters.get("seed"), + "num_images": tool_parameters.get("num_images", 1), + "safety_tolerance": tool_parameters.get("safety_tolerance", "2"), + "sync_mode": tool_parameters.get("sync_mode", False), + } + + url = "https://fal.run/fal-ai/flux-pro/new" + + response = requests.post(url, json=payload, headers=headers) + + if response.status_code != 200: + return self.create_text_message(f"Got Error Response: {response.text}") + + res = response.json() + result = [self.create_json_message(res)] + + for image_info in res.get("images", []): + image_url = image_info.get("url") + if image_url: + result.append(self.create_image_message(image=image_url, save_as=self.VariableKey.IMAGE.value)) + + return result diff --git a/api/core/tools/provider/builtin/fal/tools/flux_1_pro_new.yaml b/api/core/tools/provider/builtin/fal/tools/flux_1_pro_new.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6f8dbb3a54e9ec0eaf0665e9f1007ba37120755c --- /dev/null +++ b/api/core/tools/provider/builtin/fal/tools/flux_1_pro_new.yaml @@ -0,0 +1,164 @@ +identity: + name: flux_1_pro_new + author: Kalo Chin + label: + en_US: FLUX.1 [pro] new + zh_Hans: FLUX.1 [pro] new + icon: icon.svg +description: + human: + en_US: FLUX.1 [pro] new is an accelerated version of FLUX.1 [pro], maintaining professional-grade image quality while delivering significantly faster generation speeds. + zh_Hans: FLUX.1 [pro] new 是 FLUX.1 [pro] 的加速版本,在保持专业级图像质量的同时,大大提高了生成速度。 + llm: This tool generates images from prompts using FAL's FLUX.1 [pro] new model. +parameters: + - name: prompt + type: string + required: true + label: + en_US: Prompt + zh_Hans: 提示词 + human_description: + en_US: The text prompt used to generate the image. + zh_Hans: 用于生成图像的文本提示。 + llm_description: This prompt text will be used to generate the image. + form: llm + - name: image_size + type: select + required: false + options: + - value: square_hd + label: + en_US: Square HD + zh_Hans: 正方形高清 + - value: square + label: + en_US: Square + zh_Hans: 正方形 + - value: portrait_4_3 + label: + en_US: Portrait 4:3 + zh_Hans: 竖屏 4:3 + - value: portrait_16_9 + label: + en_US: Portrait 16:9 + zh_Hans: 竖屏 16:9 + - value: landscape_4_3 + label: + en_US: Landscape 4:3 + zh_Hans: 横屏 4:3 + - value: landscape_16_9 + label: + en_US: Landscape 16:9 + zh_Hans: 横屏 16:9 + default: landscape_4_3 + label: + en_US: Image Size + zh_Hans: 图像尺寸 + human_description: + en_US: The size of the generated image. + zh_Hans: 生成图像的尺寸。 + form: form + - name: num_images + type: number + required: false + default: 1 + min: 1 + max: 1 + label: + en_US: Number of Images + zh_Hans: 图像数量 + human_description: + en_US: The number of images to generate. + zh_Hans: 要生成的图像数量。 + form: form + - name: num_inference_steps + type: number + required: false + default: 28 + min: 1 + max: 50 + label: + en_US: Num Inference Steps + zh_Hans: 推理步数 + human_description: + en_US: The number of inference steps to perform. More steps produce higher quality but take longer. + zh_Hans: 执行的推理步数。步数越多,质量越高,但所需时间也更长。 + form: form + - name: guidance_scale + type: number + required: false + default: 3.5 + min: 0 + max: 20 + label: + en_US: Guidance Scale + zh_Hans: 指导强度 + human_description: + en_US: How closely the model should follow the prompt. + zh_Hans: 模型对提示词的遵循程度。 + form: form + - name: safety_tolerance + type: select + required: false + options: + - value: "1" + label: + en_US: "1 (Most strict)" + zh_Hans: "1(最严格)" + - value: "2" + label: + en_US: "2" + zh_Hans: "2" + - value: "3" + label: + en_US: "3" + zh_Hans: "3" + - value: "4" + label: + en_US: "4" + zh_Hans: "4" + - value: "5" + label: + en_US: "5" + zh_Hans: "5" + - value: "6" + label: + en_US: "6 (Most permissive)" + zh_Hans: "6(最宽松)" + default: "2" + label: + en_US: Safety Tolerance + zh_Hans: 安全容忍度 + human_description: + en_US: > + The safety tolerance level for the generated image. 1 being the most strict and 5 being the most permissive. + zh_Hans: > + 生成图像的安全容忍级别。1 是最严格,6 是最宽松。 + form: form + - name: seed + type: number + required: false + min: 0 + max: 9999999999 + label: + en_US: Seed + zh_Hans: 种子 + human_description: + en_US: The same seed and prompt can produce similar images. + zh_Hans: 相同的种子和提示词可以生成相似的图像。 + form: form + - name: sync_mode + type: boolean + required: false + default: false + label: + en_US: Sync Mode + zh_Hans: 同步模式 + human_description: + en_US: > + If set to true, the function will wait for the image to be generated and uploaded before returning the response. + This will increase the latency but allows you to get the image directly in the response without going through the CDN. + zh_Hans: > + 如果设置为 true,函数将在生成并上传图像后才返回响应。 + 这将增加延迟,但允许您直接在响应中获取图像,而无需通过 CDN。 + form: form diff --git a/api/core/tools/provider/builtin/fal/tools/wizper.py b/api/core/tools/provider/builtin/fal/tools/wizper.py new file mode 100644 index 0000000000000000000000000000000000000000..ba05a6207330b58497df5c42367cfde3feb8e574 --- /dev/null +++ b/api/core/tools/provider/builtin/fal/tools/wizper.py @@ -0,0 +1,56 @@ +import io +import os +from typing import Any + +import fal_client + +from core.file.enums import FileAttribute, FileType +from core.file.file_manager import download, get_attr +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class WizperTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + audio_file = tool_parameters.get("audio_file") + task = tool_parameters.get("task", "transcribe") + language = tool_parameters.get("language", "en") + chunk_level = tool_parameters.get("chunk_level", "segment") + version = tool_parameters.get("version", "3") + + if audio_file.type != FileType.AUDIO: + return self.create_text_message("Not a valid audio file.") + + api_key = self.runtime.credentials["fal_api_key"] + + os.environ["FAL_KEY"] = api_key + + audio_binary = io.BytesIO(download(audio_file)) + mime_type = get_attr(file=audio_file, attr=FileAttribute.MIME_TYPE) + file_data = audio_binary.getvalue() + + try: + audio_url = fal_client.upload(file_data, mime_type) + except Exception as e: + return self.create_text_message(f"Error uploading audio file: {str(e)}") + + arguments = { + "audio_url": audio_url, + "task": task, + "language": language, + "chunk_level": chunk_level, + "version": version, + } + + result = fal_client.subscribe( + "fal-ai/wizper", + arguments=arguments, + with_logs=False, + ) + + json_message = self.create_json_message(result) + + text = result.get("text", "") + text_message = self.create_text_message(text) + + return [json_message, text_message] diff --git a/api/core/tools/provider/builtin/fal/tools/wizper.yaml b/api/core/tools/provider/builtin/fal/tools/wizper.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5742efcc1b40025de746c2d702e151b20fb2ee6f --- /dev/null +++ b/api/core/tools/provider/builtin/fal/tools/wizper.yaml @@ -0,0 +1,489 @@ +identity: + name: wizper + author: Kalo Chin + label: + en_US: Wizper + zh_Hans: Wizper +description: + human: + en_US: Transcribe an audio file using the Whisper model. + zh_Hans: 使用 Whisper 模型转录音频文件。 + llm: Transcribe an audio file using the Whisper model. +parameters: + - name: audio_file + type: file + required: true + label: + en_US: Audio File + zh_Hans: 音频文件 + human_description: + en_US: "Upload an audio file to transcribe. Supports mp3, mp4, mpeg, mpga, m4a, wav, or webm formats." + zh_Hans: "上传要转录的音频文件。支持 mp3、mp4、mpeg、mpga、m4a、wav 或 webm 格式。" + llm_description: "Audio file to transcribe. Supported formats: mp3, mp4, mpeg, mpga, m4a, wav, or webm." + form: llm + - name: task + type: select + required: true + label: + en_US: Task + zh_Hans: 任务 + human_description: + en_US: "Choose whether to transcribe the audio in its original language or translate it to English" + zh_Hans: "选择是以原始语言转录音频还是将其翻译成英语" + llm_description: "Task to perform on the audio file. Either transcribe or translate. Default value: 'transcribe'. If 'translate' is selected as the task, the audio will be translated to English, regardless of the language selected." + form: form + default: transcribe + options: + - value: transcribe + label: + en_US: Transcribe + zh_Hans: 转录 + - value: translate + label: + en_US: Translate + zh_Hans: 翻译 + - name: language + type: select + required: true + label: + en_US: Language + zh_Hans: 语言 + human_description: + en_US: "Select the primary language spoken in the audio file" + zh_Hans: "选择音频文件中使用的主要语言" + llm_description: "Language of the audio file." + form: form + default: en + options: + - value: af + label: + en_US: Afrikaans + zh_Hans: 南非语 + - value: am + label: + en_US: Amharic + zh_Hans: 阿姆哈拉语 + - value: ar + label: + en_US: Arabic + zh_Hans: 阿拉伯语 + - value: as + label: + en_US: Assamese + zh_Hans: 阿萨姆语 + - value: az + label: + en_US: Azerbaijani + zh_Hans: 阿塞拜疆语 + - value: ba + label: + en_US: Bashkir + zh_Hans: 巴什基尔语 + - value: be + label: + en_US: Belarusian + zh_Hans: 白俄罗斯语 + - value: bg + label: + en_US: Bulgarian + zh_Hans: 保加利亚语 + - value: bn + label: + en_US: Bengali + zh_Hans: 孟加拉语 + - value: bo + label: + en_US: Tibetan + zh_Hans: 藏语 + - value: br + label: + en_US: Breton + zh_Hans: 布列塔尼语 + - value: bs + label: + en_US: Bosnian + zh_Hans: 波斯尼亚语 + - value: ca + label: + en_US: Catalan + zh_Hans: 加泰罗尼亚语 + - value: cs + label: + en_US: Czech + zh_Hans: 捷克语 + - value: cy + label: + en_US: Welsh + zh_Hans: 威尔士语 + - value: da + label: + en_US: Danish + zh_Hans: 丹麦语 + - value: de + label: + en_US: German + zh_Hans: 德语 + - value: el + label: + en_US: Greek + zh_Hans: 希腊语 + - value: en + label: + en_US: English + zh_Hans: 英语 + - value: es + label: + en_US: Spanish + zh_Hans: 西班牙语 + - value: et + label: + en_US: Estonian + zh_Hans: 爱沙尼亚语 + - value: eu + label: + en_US: Basque + zh_Hans: 巴斯克语 + - value: fa + label: + en_US: Persian + zh_Hans: 波斯语 + - value: fi + label: + en_US: Finnish + zh_Hans: 芬兰语 + - value: fo + label: + en_US: Faroese + zh_Hans: 法罗语 + - value: fr + label: + en_US: French + zh_Hans: 法语 + - value: gl + label: + en_US: Galician + zh_Hans: 加利西亚语 + - value: gu + label: + en_US: Gujarati + zh_Hans: 古吉拉特语 + - value: ha + label: + en_US: Hausa + zh_Hans: 毫萨语 + - value: haw + label: + en_US: Hawaiian + zh_Hans: 夏威夷语 + - value: he + label: + en_US: Hebrew + zh_Hans: 希伯来语 + - value: hi + label: + en_US: Hindi + zh_Hans: 印地语 + - value: hr + label: + en_US: Croatian + zh_Hans: 克罗地亚语 + - value: ht + label: + en_US: Haitian Creole + zh_Hans: 海地克里奥尔语 + - value: hu + label: + en_US: Hungarian + zh_Hans: 匈牙利语 + - value: hy + label: + en_US: Armenian + zh_Hans: 亚美尼亚语 + - value: id + label: + en_US: Indonesian + zh_Hans: 印度尼西亚语 + - value: is + label: + en_US: Icelandic + zh_Hans: 冰岛语 + - value: it + label: + en_US: Italian + zh_Hans: 意大利语 + - value: ja + label: + en_US: Japanese + zh_Hans: 日语 + - value: jw + label: + en_US: Javanese + zh_Hans: 爪哇语 + - value: ka + label: + en_US: Georgian + zh_Hans: 格鲁吉亚语 + - value: kk + label: + en_US: Kazakh + zh_Hans: 哈萨克语 + - value: km + label: + en_US: Khmer + zh_Hans: 高棉语 + - value: kn + label: + en_US: Kannada + zh_Hans: 卡纳达语 + - value: ko + label: + en_US: Korean + zh_Hans: 韩语 + - value: la + label: + en_US: Latin + zh_Hans: 拉丁语 + - value: lb + label: + en_US: Luxembourgish + zh_Hans: 卢森堡语 + - value: ln + label: + en_US: Lingala + zh_Hans: 林加拉语 + - value: lo + label: + en_US: Lao + zh_Hans: 老挝语 + - value: lt + label: + en_US: Lithuanian + zh_Hans: 立陶宛语 + - value: lv + label: + en_US: Latvian + zh_Hans: 拉脱维亚语 + - value: mg + label: + en_US: Malagasy + zh_Hans: 马尔加什语 + - value: mi + label: + en_US: Maori + zh_Hans: 毛利语 + - value: mk + label: + en_US: Macedonian + zh_Hans: 马其顿语 + - value: ml + label: + en_US: Malayalam + zh_Hans: 马拉雅拉姆语 + - value: mn + label: + en_US: Mongolian + zh_Hans: 蒙古语 + - value: mr + label: + en_US: Marathi + zh_Hans: 马拉地语 + - value: ms + label: + en_US: Malay + zh_Hans: 马来语 + - value: mt + label: + en_US: Maltese + zh_Hans: 马耳他语 + - value: my + label: + en_US: Burmese + zh_Hans: 缅甸语 + - value: ne + label: + en_US: Nepali + zh_Hans: 尼泊尔语 + - value: nl + label: + en_US: Dutch + zh_Hans: 荷兰语 + - value: nn + label: + en_US: Norwegian Nynorsk + zh_Hans: 新挪威语 + - value: no + label: + en_US: Norwegian + zh_Hans: 挪威语 + - value: oc + label: + en_US: Occitan + zh_Hans: 奥克语 + - value: pa + label: + en_US: Punjabi + zh_Hans: 旁遮普语 + - value: pl + label: + en_US: Polish + zh_Hans: 波兰语 + - value: ps + label: + en_US: Pashto + zh_Hans: 普什图语 + - value: pt + label: + en_US: Portuguese + zh_Hans: 葡萄牙语 + - value: ro + label: + en_US: Romanian + zh_Hans: 罗马尼亚语 + - value: ru + label: + en_US: Russian + zh_Hans: 俄语 + - value: sa + label: + en_US: Sanskrit + zh_Hans: 梵语 + - value: sd + label: + en_US: Sindhi + zh_Hans: 信德语 + - value: si + label: + en_US: Sinhala + zh_Hans: 僧伽罗语 + - value: sk + label: + en_US: Slovak + zh_Hans: 斯洛伐克语 + - value: sl + label: + en_US: Slovenian + zh_Hans: 斯洛文尼亚语 + - value: sn + label: + en_US: Shona + zh_Hans: 修纳语 + - value: so + label: + en_US: Somali + zh_Hans: 索马里语 + - value: sq + label: + en_US: Albanian + zh_Hans: 阿尔巴尼亚语 + - value: sr + label: + en_US: Serbian + zh_Hans: 塞尔维亚语 + - value: su + label: + en_US: Sundanese + zh_Hans: 巽他语 + - value: sv + label: + en_US: Swedish + zh_Hans: 瑞典语 + - value: sw + label: + en_US: Swahili + zh_Hans: 斯瓦希里语 + - value: ta + label: + en_US: Tamil + zh_Hans: 泰米尔语 + - value: te + label: + en_US: Telugu + zh_Hans: 泰卢固语 + - value: tg + label: + en_US: Tajik + zh_Hans: 塔吉克语 + - value: th + label: + en_US: Thai + zh_Hans: 泰语 + - value: tk + label: + en_US: Turkmen + zh_Hans: 土库曼语 + - value: tl + label: + en_US: Tagalog + zh_Hans: 他加禄语 + - value: tr + label: + en_US: Turkish + zh_Hans: 土耳其语 + - value: tt + label: + en_US: Tatar + zh_Hans: 鞑靼语 + - value: uk + label: + en_US: Ukrainian + zh_Hans: 乌克兰语 + - value: ur + label: + en_US: Urdu + zh_Hans: 乌尔都语 + - value: uz + label: + en_US: Uzbek + zh_Hans: 乌兹别克语 + - value: vi + label: + en_US: Vietnamese + zh_Hans: 越南语 + - value: yi + label: + en_US: Yiddish + zh_Hans: 意第绪语 + - value: yo + label: + en_US: Yoruba + zh_Hans: 约鲁巴语 + - value: yue + label: + en_US: Cantonese + zh_Hans: 粤语 + - value: zh + label: + en_US: Chinese + zh_Hans: 中文 + - name: chunk_level + type: select + label: + en_US: Chunk Level + zh_Hans: 分块级别 + human_description: + en_US: "Choose how the transcription should be divided into chunks" + zh_Hans: "选择如何将转录内容分成块" + llm_description: "Level of the chunks to return." + form: form + default: segment + options: + - value: segment + label: + en_US: Segment + zh_Hans: 段 + - name: version + type: select + label: + en_US: Version + zh_Hans: 版本 + human_description: + en_US: "Select which version of the Whisper large model to use" + zh_Hans: "选择要使用的 Whisper large 模型版本" + llm_description: "Version of the model to use. All of the models are the Whisper large variant." + form: form + default: "3" + options: + - value: "3" + label: + en_US: Version 3 + zh_Hans: 版本 3 diff --git a/api/core/tools/provider/builtin/feishu/_assets/icon.svg b/api/core/tools/provider/builtin/feishu/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..bf3c202abf3ff657734f38d0aef67454cc7e2f9b --- /dev/null +++ b/api/core/tools/provider/builtin/feishu/_assets/icon.svg @@ -0,0 +1 @@ + diff --git a/api/core/tools/provider/builtin/feishu/feishu.py b/api/core/tools/provider/builtin/feishu/feishu.py new file mode 100644 index 0000000000000000000000000000000000000000..72a9333619988d425630737496315fb68214148b --- /dev/null +++ b/api/core/tools/provider/builtin/feishu/feishu.py @@ -0,0 +1,7 @@ +from core.tools.provider.builtin.feishu.tools.feishu_group_bot import FeishuGroupBotTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class FeishuProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + FeishuGroupBotTool() diff --git a/api/core/tools/provider/builtin/feishu/feishu.yaml b/api/core/tools/provider/builtin/feishu/feishu.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a029c7edb8853bbb1f98588c6f1589184f21744c --- /dev/null +++ b/api/core/tools/provider/builtin/feishu/feishu.yaml @@ -0,0 +1,16 @@ +identity: + author: Arkii Sun + name: feishu + label: + en_US: Feishu + zh_Hans: 飞书 + pt_BR: Feishu + description: + en_US: Feishu group bot + zh_Hans: 飞书群机器人 + pt_BR: Feishu group bot + icon: icon.svg + tags: + - social + - productivity +credentials_for_provider: diff --git a/api/core/tools/provider/builtin/feishu/tools/feishu_group_bot.py b/api/core/tools/provider/builtin/feishu/tools/feishu_group_bot.py new file mode 100644 index 0000000000000000000000000000000000000000..e82da8ca534b96729b6eda414009e15c25b64835 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu/tools/feishu_group_bot.py @@ -0,0 +1,51 @@ +from typing import Any, Union + +import httpx + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.uuid_utils import is_valid_uuid + + +class FeishuGroupBotTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + API document: https://open.feishu.cn/document/client-docs/bot-v3/add-custom-bot + """ + + url = "https://open.feishu.cn/open-apis/bot/v2/hook" + + content = tool_parameters.get("content", "") + if not content: + return self.create_text_message("Invalid parameter content") + + hook_key = tool_parameters.get("hook_key", "") + if not is_valid_uuid(hook_key): + return self.create_text_message(f"Invalid parameter hook_key ${hook_key}, not a valid UUID") + + msg_type = "text" + api_url = f"{url}/{hook_key}" + headers = { + "Content-Type": "application/json", + } + params = {} + payload = { + "msg_type": msg_type, + "content": { + "text": content, + }, + } + + try: + res = httpx.post(api_url, headers=headers, params=params, json=payload) + if res.is_success: + return self.create_text_message("Text message sent successfully") + else: + return self.create_text_message( + f"Failed to send the text message, status code: {res.status_code}, response: {res.text}" + ) + except Exception as e: + return self.create_text_message("Failed to send message to group chat bot. {}".format(e)) diff --git a/api/core/tools/provider/builtin/feishu/tools/feishu_group_bot.yaml b/api/core/tools/provider/builtin/feishu/tools/feishu_group_bot.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6c3f084e4dafe393f9422438dd868ce0037348fd --- /dev/null +++ b/api/core/tools/provider/builtin/feishu/tools/feishu_group_bot.yaml @@ -0,0 +1,40 @@ +identity: + name: feishu_group_bot + author: Arkii Sun + label: + en_US: Send Group Message + zh_Hans: 发送群消息 + pt_BR: Send Group Message + icon: icon.png +description: + human: + en_US: Sending a group message on Feishu via the webhook of group bot + zh_Hans: 通过飞书的群机器人webhook发送群消息 + pt_BR: Sending a group message on Feishu via the webhook of group bot + llm: A tool for sending messages to a chat group on Feishu(飞书) . +parameters: + - name: hook_key + type: secret-input + required: true + label: + en_US: Feishu Group bot webhook key + zh_Hans: 群机器人webhook的key + pt_BR: Feishu Group bot webhook key + human_description: + en_US: Feishu Group bot webhook key + zh_Hans: 群机器人webhook的key + pt_BR: Feishu Group bot webhook key + form: form + - name: content + type: string + required: true + label: + en_US: content + zh_Hans: 消息内容 + pt_BR: content + human_description: + en_US: Content to sent to the group. + zh_Hans: 群消息文本 + pt_BR: Content to sent to the group. + llm_description: Content of the message + form: llm diff --git a/api/core/tools/provider/builtin/feishu_base/_assets/icon.png b/api/core/tools/provider/builtin/feishu_base/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..787427e7218058678986333768b33a0aafd1eb58 Binary files /dev/null and b/api/core/tools/provider/builtin/feishu_base/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/feishu_base/feishu_base.py b/api/core/tools/provider/builtin/feishu_base/feishu_base.py new file mode 100644 index 0000000000000000000000000000000000000000..f301ec5355d48f38d454d93999cb04aa0061a1b3 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/feishu_base.py @@ -0,0 +1,7 @@ +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController +from core.tools.utils.feishu_api_utils import auth + + +class FeishuBaseProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + auth(credentials) diff --git a/api/core/tools/provider/builtin/feishu_base/feishu_base.yaml b/api/core/tools/provider/builtin/feishu_base/feishu_base.yaml new file mode 100644 index 0000000000000000000000000000000000000000..456dd8c88fc34829e7bec484e1d4ff0088890a42 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/feishu_base.yaml @@ -0,0 +1,36 @@ +identity: + author: Doug Lea + name: feishu_base + label: + en_US: Feishu Base + zh_Hans: 飞书多维表格 + description: + en_US: | + Feishu base, requires the following permissions: bitable:app. + zh_Hans: | + 飞书多维表格,需要开通以下权限: bitable:app。 + icon: icon.png + tags: + - social + - productivity +credentials_for_provider: + app_id: + type: text-input + required: true + label: + en_US: APP ID + placeholder: + en_US: Please input your feishu app id + zh_Hans: 请输入你的飞书 app id + help: + en_US: Get your app_id and app_secret from Feishu + zh_Hans: 从飞书获取您的 app_id 和 app_secret + url: https://open.larkoffice.com/app + app_secret: + type: secret-input + required: true + label: + en_US: APP Secret + placeholder: + en_US: Please input your app secret + zh_Hans: 请输入你的飞书 app secret diff --git a/api/core/tools/provider/builtin/feishu_base/tools/add_records.py b/api/core/tools/provider/builtin/feishu_base/tools/add_records.py new file mode 100644 index 0000000000000000000000000000000000000000..905f8b78806d0541e884897ae33a978c2d445da5 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/tools/add_records.py @@ -0,0 +1,21 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class AddRecordsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + app_token = tool_parameters.get("app_token") + table_id = tool_parameters.get("table_id") + table_name = tool_parameters.get("table_name") + records = tool_parameters.get("records") + user_id_type = tool_parameters.get("user_id_type", "open_id") + + res = client.add_records(app_token, table_id, table_name, records, user_id_type) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_base/tools/add_records.yaml b/api/core/tools/provider/builtin/feishu_base/tools/add_records.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f2a93490dc0c3103b9fe4de29f54d4bdd6db5bdc --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/tools/add_records.yaml @@ -0,0 +1,91 @@ +identity: + name: add_records + author: Doug Lea + label: + en_US: Add Records + zh_Hans: 新增多条记录 +description: + human: + en_US: Add Multiple Records to Multidimensional Table + zh_Hans: 在多维表格数据表中新增多条记录 + llm: A tool for adding multiple records to a multidimensional table. (在多维表格数据表中新增多条记录) +parameters: + - name: app_token + type: string + required: true + label: + en_US: app_token + zh_Hans: app_token + human_description: + en_US: Unique identifier for the multidimensional table, supports inputting document URL. + zh_Hans: 多维表格的唯一标识符,支持输入文档 URL。 + llm_description: 多维表格的唯一标识符,支持输入文档 URL。 + form: llm + + - name: table_id + type: string + required: false + label: + en_US: table_id + zh_Hans: table_id + human_description: + en_US: Unique identifier for the multidimensional table data, either table_id or table_name must be provided, cannot be empty simultaneously. + zh_Hans: 多维表格数据表的唯一标识符,table_id 和 table_name 至少需要提供一个,不能同时为空。 + llm_description: 多维表格数据表的唯一标识符,table_id 和 table_name 至少需要提供一个,不能同时为空。 + form: llm + + - name: table_name + type: string + required: false + label: + en_US: table_name + zh_Hans: table_name + human_description: + en_US: Name of the multidimensional table data, either table_name or table_id must be provided, cannot be empty simultaneously. + zh_Hans: 多维表格数据表的名称,table_name 和 table_id 至少需要提供一个,不能同时为空。 + llm_description: 多维表格数据表的名称,table_name 和 table_id 至少需要提供一个,不能同时为空。 + form: llm + + - name: records + type: string + required: true + label: + en_US: records + zh_Hans: 记录列表 + human_description: + en_US: | + List of records to be added in this request. Example value: [{"multi-line-text":"text content","single_select":"option 1","date":1674206443000}] + For supported field types, refer to the integration guide (https://open.larkoffice.com/document/server-docs/docs/bitable-v1/notification). For data structures of different field types, refer to the data structure overview (https://open.larkoffice.com/document/server-docs/docs/bitable-v1/bitable-structure). + zh_Hans: | + 本次请求将要新增的记录列表,示例值:[{"多行文本":"文本内容","单选":"选项 1","日期":1674206443000}]。 + 当前接口支持的字段类型请参考接入指南(https://open.larkoffice.com/document/server-docs/docs/bitable-v1/notification),不同类型字段的数据结构请参考数据结构概述(https://open.larkoffice.com/document/server-docs/docs/bitable-v1/bitable-structure)。 + llm_description: | + 本次请求将要新增的记录列表,示例值:[{"多行文本":"文本内容","单选":"选项 1","日期":1674206443000}]。 + 当前接口支持的字段类型请参考接入指南(https://open.larkoffice.com/document/server-docs/docs/bitable-v1/notification),不同类型字段的数据结构请参考数据结构概述(https://open.larkoffice.com/document/server-docs/docs/bitable-v1/bitable-structure)。 + form: llm + + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form diff --git a/api/core/tools/provider/builtin/feishu_base/tools/create_base.py b/api/core/tools/provider/builtin/feishu_base/tools/create_base.py new file mode 100644 index 0000000000000000000000000000000000000000..f074acc5ff709e6b9ab92c398aea46c824d03fb7 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/tools/create_base.py @@ -0,0 +1,18 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class CreateBaseTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + name = tool_parameters.get("name") + folder_token = tool_parameters.get("folder_token") + + res = client.create_base(name, folder_token) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_base/tools/create_base.yaml b/api/core/tools/provider/builtin/feishu_base/tools/create_base.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3ec91a90e7f0b6737180efa11f7d0241b8c270fc --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/tools/create_base.yaml @@ -0,0 +1,42 @@ +identity: + name: create_base + author: Doug Lea + label: + en_US: Create Base + zh_Hans: 创建多维表格 +description: + human: + en_US: Create Multidimensional Table in Specified Directory + zh_Hans: 在指定目录下创建多维表格 + llm: A tool for creating a multidimensional table in a specified directory. (在指定目录下创建多维表格) +parameters: + - name: name + type: string + required: false + label: + en_US: name + zh_Hans: 多维表格 App 名字 + human_description: + en_US: | + Name of the multidimensional table App. Example value: "A new multidimensional table". + zh_Hans: 多维表格 App 名字,示例值:"一篇新的多维表格"。 + llm_description: 多维表格 App 名字,示例值:"一篇新的多维表格"。 + form: llm + + - name: folder_token + type: string + required: false + label: + en_US: folder_token + zh_Hans: 多维表格 App 归属文件夹 + human_description: + en_US: | + Folder where the multidimensional table App belongs. Default is empty, meaning the table will be created in the root directory of the cloud space. Example values: Fa3sfoAgDlMZCcdcJy1cDFg8nJc or https://svi136aogf123.feishu.cn/drive/folder/Fa3sfoAgDlMZCcdcJy1cDFg8nJc. + The folder_token must be an existing folder and supports inputting folder token or folder URL. + zh_Hans: | + 多维表格 App 归属文件夹。默认为空,表示多维表格将被创建在云空间根目录。示例值: Fa3sfoAgDlMZCcdcJy1cDFg8nJc 或者 https://svi136aogf123.feishu.cn/drive/folder/Fa3sfoAgDlMZCcdcJy1cDFg8nJc。 + folder_token 必须是已存在的文件夹,支持输入文件夹 token 或者文件夹 URL。 + llm_description: | + 多维表格 App 归属文件夹。默认为空,表示多维表格将被创建在云空间根目录。示例值: Fa3sfoAgDlMZCcdcJy1cDFg8nJc 或者 https://svi136aogf123.feishu.cn/drive/folder/Fa3sfoAgDlMZCcdcJy1cDFg8nJc。 + folder_token 必须是已存在的文件夹,支持输入文件夹 token 或者文件夹 URL。 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_base/tools/create_table.py b/api/core/tools/provider/builtin/feishu_base/tools/create_table.py new file mode 100644 index 0000000000000000000000000000000000000000..81f2617545969bbbc32a439dd2997549b733e230 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/tools/create_table.py @@ -0,0 +1,20 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class CreateTableTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + app_token = tool_parameters.get("app_token") + table_name = tool_parameters.get("table_name") + default_view_name = tool_parameters.get("default_view_name") + fields = tool_parameters.get("fields") + + res = client.create_table(app_token, table_name, default_view_name, fields) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_base/tools/create_table.yaml b/api/core/tools/provider/builtin/feishu_base/tools/create_table.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8b1007b9a531663b87846dfcdd3f075b81929420 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/tools/create_table.yaml @@ -0,0 +1,61 @@ +identity: + name: create_table + author: Doug Lea + label: + en_US: Create Table + zh_Hans: 新增数据表 +description: + human: + en_US: Add a Data Table to Multidimensional Table + zh_Hans: 在多维表格中新增一个数据表 + llm: A tool for adding a data table to a multidimensional table. (在多维表格中新增一个数据表) +parameters: + - name: app_token + type: string + required: true + label: + en_US: app_token + zh_Hans: app_token + human_description: + en_US: Unique identifier for the multidimensional table, supports inputting document URL. + zh_Hans: 多维表格的唯一标识符,支持输入文档 URL。 + llm_description: 多维表格的唯一标识符,支持输入文档 URL。 + form: llm + + - name: table_name + type: string + required: true + label: + en_US: Table Name + zh_Hans: 数据表名称 + human_description: + en_US: | + The name of the data table, length range: 1 character to 100 characters. + zh_Hans: 数据表名称,长度范围:1 字符 ~ 100 字符。 + llm_description: 数据表名称,长度范围:1 字符 ~ 100 字符。 + form: llm + + - name: default_view_name + type: string + required: false + label: + en_US: Default View Name + zh_Hans: 默认表格视图的名称 + human_description: + en_US: The name of the default table view, defaults to "Table" if not filled. + zh_Hans: 默认表格视图的名称,不填则默认为"表格"。 + llm_description: 默认表格视图的名称,不填则默认为"表格"。 + form: llm + + - name: fields + type: string + required: true + label: + en_US: Initial Fields + zh_Hans: 初始字段 + human_description: + en_US: | + Initial fields of the data table, format: [ { "field_name": "Multi-line Text","type": 1 },{ "field_name": "Number","type": 2 },{ "field_name": "Single Select","type": 3 },{ "field_name": "Multiple Select","type": 4 },{ "field_name": "Date","type": 5 } ]. For field details, refer to: https://open.larkoffice.com/document/server-docs/docs/bitable-v1/app-table-field/guide + zh_Hans: 数据表的初始字段,格式为:[{"field_name":"多行文本","type":1},{"field_name":"数字","type":2},{"field_name":"单选","type":3},{"field_name":"多选","type":4},{"field_name":"日期","type":5}]。字段详情参考:https://open.larkoffice.com/document/server-docs/docs/bitable-v1/app-table-field/guide + llm_description: 数据表的初始字段,格式为:[{"field_name":"多行文本","type":1},{"field_name":"数字","type":2},{"field_name":"单选","type":3},{"field_name":"多选","type":4},{"field_name":"日期","type":5}]。字段详情参考:https://open.larkoffice.com/document/server-docs/docs/bitable-v1/app-table-field/guide + form: llm diff --git a/api/core/tools/provider/builtin/feishu_base/tools/delete_records.py b/api/core/tools/provider/builtin/feishu_base/tools/delete_records.py new file mode 100644 index 0000000000000000000000000000000000000000..c896a2c81b97f860da4ec8322b4155027b5c9003 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/tools/delete_records.py @@ -0,0 +1,20 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class DeleteRecordsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + app_token = tool_parameters.get("app_token") + table_id = tool_parameters.get("table_id") + table_name = tool_parameters.get("table_name") + record_ids = tool_parameters.get("record_ids") + + res = client.delete_records(app_token, table_id, table_name, record_ids) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_base/tools/delete_records.yaml b/api/core/tools/provider/builtin/feishu_base/tools/delete_records.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c30ebd630ce9d835a78fa77724cecf16acfe5dbe --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/tools/delete_records.yaml @@ -0,0 +1,86 @@ +identity: + name: delete_records + author: Doug Lea + label: + en_US: Delete Records + zh_Hans: 删除多条记录 +description: + human: + en_US: Delete Multiple Records from Multidimensional Table + zh_Hans: 删除多维表格数据表中的多条记录 + llm: A tool for deleting multiple records from a multidimensional table. (删除多维表格数据表中的多条记录) +parameters: + - name: app_token + type: string + required: true + label: + en_US: app_token + zh_Hans: app_token + human_description: + en_US: Unique identifier for the multidimensional table, supports inputting document URL. + zh_Hans: 多维表格的唯一标识符,支持输入文档 URL。 + llm_description: 多维表格的唯一标识符,支持输入文档 URL。 + form: llm + + - name: table_id + type: string + required: false + label: + en_US: table_id + zh_Hans: table_id + human_description: + en_US: Unique identifier for the multidimensional table data, either table_id or table_name must be provided, cannot be empty simultaneously. + zh_Hans: 多维表格数据表的唯一标识符,table_id 和 table_name 至少需要提供一个,不能同时为空。 + llm_description: 多维表格数据表的唯一标识符,table_id 和 table_name 至少需要提供一个,不能同时为空。 + form: llm + + - name: table_name + type: string + required: false + label: + en_US: table_name + zh_Hans: table_name + human_description: + en_US: Name of the multidimensional table data, either table_name or table_id must be provided, cannot be empty simultaneously. + zh_Hans: 多维表格数据表的名称,table_name 和 table_id 至少需要提供一个,不能同时为空。 + llm_description: 多维表格数据表的名称,table_name 和 table_id 至少需要提供一个,不能同时为空。 + form: llm + + - name: record_ids + type: string + required: true + label: + en_US: Record IDs + zh_Hans: 记录 ID 列表 + human_description: + en_US: | + List of IDs for the records to be deleted, example value: ["recwNXzPQv"]. + zh_Hans: 删除的多条记录 ID 列表,示例值:["recwNXzPQv"]。 + llm_description: 删除的多条记录 ID 列表,示例值:["recwNXzPQv"]。 + form: llm + + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form diff --git a/api/core/tools/provider/builtin/feishu_base/tools/delete_tables.py b/api/core/tools/provider/builtin/feishu_base/tools/delete_tables.py new file mode 100644 index 0000000000000000000000000000000000000000..f732a16da6f69794e7e7854a8fc5321652f3ade1 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/tools/delete_tables.py @@ -0,0 +1,19 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class DeleteTablesTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + app_token = tool_parameters.get("app_token") + table_ids = tool_parameters.get("table_ids") + table_names = tool_parameters.get("table_names") + + res = client.delete_tables(app_token, table_ids, table_names) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_base/tools/delete_tables.yaml b/api/core/tools/provider/builtin/feishu_base/tools/delete_tables.yaml new file mode 100644 index 0000000000000000000000000000000000000000..498126eae53302d088f275d9f3fc71c9b6cff378 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/tools/delete_tables.yaml @@ -0,0 +1,49 @@ +identity: + name: delete_tables + author: Doug Lea + label: + en_US: Delete Tables + zh_Hans: 删除数据表 +description: + human: + en_US: Batch Delete Data Tables from Multidimensional Table + zh_Hans: 批量删除多维表格中的数据表 + llm: A tool for batch deleting data tables from a multidimensional table. (批量删除多维表格中的数据表) +parameters: + - name: app_token + type: string + required: true + label: + en_US: app_token + zh_Hans: app_token + human_description: + en_US: Unique identifier for the multidimensional table, supports inputting document URL. + zh_Hans: 多维表格的唯一标识符,支持输入文档 URL。 + llm_description: 多维表格的唯一标识符,支持输入文档 URL。 + form: llm + + - name: table_ids + type: string + required: false + label: + en_US: Table IDs + zh_Hans: 数据表 ID + human_description: + en_US: | + IDs of the tables to be deleted. Each operation supports deleting up to 50 tables. Example: ["tbl1TkhyTWDkSoZ3"]. Ensure that either table_ids or table_names is not empty. + zh_Hans: 待删除的数据表的 ID,每次操作最多支持删除 50 个数据表。示例值:["tbl1TkhyTWDkSoZ3"]。请确保 table_ids 和 table_names 至少有一个不为空。 + llm_description: 待删除的数据表的 ID,每次操作最多支持删除 50 个数据表。示例值:["tbl1TkhyTWDkSoZ3"]。请确保 table_ids 和 table_names 至少有一个不为空。 + form: llm + + - name: table_names + type: string + required: false + label: + en_US: Table Names + zh_Hans: 数据表名称 + human_description: + en_US: | + Names of the tables to be deleted. Each operation supports deleting up to 50 tables. Example: ["Table1", "Table2"]. Ensure that either table_names or table_ids is not empty. + zh_Hans: 待删除的数据表的名称,每次操作最多支持删除 50 个数据表。示例值:["数据表1", "数据表2"]。请确保 table_names 和 table_ids 至少有一个不为空。 + llm_description: 待删除的数据表的名称,每次操作最多支持删除 50 个数据表。示例值:["数据表1", "数据表2"]。请确保 table_names 和 table_ids 至少有一个不为空。 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_base/tools/get_base_info.py b/api/core/tools/provider/builtin/feishu_base/tools/get_base_info.py new file mode 100644 index 0000000000000000000000000000000000000000..a74e9be288bc17573ad9ce836c4576b69e8baf4f --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/tools/get_base_info.py @@ -0,0 +1,17 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class GetBaseInfoTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + app_token = tool_parameters.get("app_token") + + res = client.get_base_info(app_token) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_base/tools/get_base_info.yaml b/api/core/tools/provider/builtin/feishu_base/tools/get_base_info.yaml new file mode 100644 index 0000000000000000000000000000000000000000..eb0e7a26c06a557b6335b82b7f46825cfabf8b5f --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/tools/get_base_info.yaml @@ -0,0 +1,23 @@ +identity: + name: get_base_info + author: Doug Lea + label: + en_US: Get Base Info + zh_Hans: 获取多维表格元数据 +description: + human: + en_US: Get Metadata Information of Specified Multidimensional Table + zh_Hans: 获取指定多维表格的元数据信息 + llm: A tool for getting metadata information of a specified multidimensional table. (获取指定多维表格的元数据信息) +parameters: + - name: app_token + type: string + required: true + label: + en_US: app_token + zh_Hans: app_token + human_description: + en_US: Unique identifier for the multidimensional table, supports inputting document URL. + zh_Hans: 多维表格的唯一标识符,支持输入文档 URL。 + llm_description: 多维表格的唯一标识符,支持输入文档 URL。 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_base/tools/list_tables.py b/api/core/tools/provider/builtin/feishu_base/tools/list_tables.py new file mode 100644 index 0000000000000000000000000000000000000000..c7768a496debce3fcb1064f56631ad84360b4fe0 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/tools/list_tables.py @@ -0,0 +1,19 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class ListTablesTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + app_token = tool_parameters.get("app_token") + page_token = tool_parameters.get("page_token") + page_size = tool_parameters.get("page_size", 20) + + res = client.list_tables(app_token, page_token, page_size) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_base/tools/list_tables.yaml b/api/core/tools/provider/builtin/feishu_base/tools/list_tables.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7571519039bd242132cb3655378be85e60461111 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/tools/list_tables.yaml @@ -0,0 +1,50 @@ +identity: + name: list_tables + author: Doug Lea + label: + en_US: List Tables + zh_Hans: 列出数据表 +description: + human: + en_US: Get All Data Tables under Multidimensional Table + zh_Hans: 获取多维表格下的所有数据表 + llm: A tool for getting all data tables under a multidimensional table. (获取多维表格下的所有数据表) +parameters: + - name: app_token + type: string + required: true + label: + en_US: app_token + zh_Hans: app_token + human_description: + en_US: Unique identifier for the multidimensional table, supports inputting document URL. + zh_Hans: 多维表格的唯一标识符,支持输入文档 URL。 + llm_description: 多维表格的唯一标识符,支持输入文档 URL。 + form: llm + + - name: page_size + type: number + required: false + default: 20 + label: + en_US: page_size + zh_Hans: 分页大小 + human_description: + en_US: | + Page size, default value: 20, maximum value: 100. + zh_Hans: 分页大小,默认值:20,最大值:100。 + llm_description: 分页大小,默认值:20,最大值:100。 + form: form + + - name: page_token + type: string + required: false + label: + en_US: page_token + zh_Hans: 分页标记 + human_description: + en_US: | + Page token, leave empty for the first request to start from the beginning; a new page_token will be returned if there are more items in the paginated query results, which can be used for the next traversal. Example value: "tblsRc9GRRXKqhvW". + zh_Hans: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。示例值:"tblsRc9GRRXKqhvW"。 + llm_description: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。示例值:"tblsRc9GRRXKqhvW"。 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_base/tools/read_records.py b/api/core/tools/provider/builtin/feishu_base/tools/read_records.py new file mode 100644 index 0000000000000000000000000000000000000000..46f3df4ff040f3b586856db1a60b5913f832f9e1 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/tools/read_records.py @@ -0,0 +1,21 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class ReadRecordsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + app_token = tool_parameters.get("app_token") + table_id = tool_parameters.get("table_id") + table_name = tool_parameters.get("table_name") + record_ids = tool_parameters.get("record_ids") + user_id_type = tool_parameters.get("user_id_type", "open_id") + + res = client.read_records(app_token, table_id, table_name, record_ids, user_id_type) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_base/tools/read_records.yaml b/api/core/tools/provider/builtin/feishu_base/tools/read_records.yaml new file mode 100644 index 0000000000000000000000000000000000000000..911e667cfc90adf5890378f50c858376f58b569d --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/tools/read_records.yaml @@ -0,0 +1,86 @@ +identity: + name: read_records + author: Doug Lea + label: + en_US: Read Records + zh_Hans: 批量获取记录 +description: + human: + en_US: Batch Retrieve Records from Multidimensional Table + zh_Hans: 批量获取多维表格数据表中的记录信息 + llm: A tool for batch retrieving records from a multidimensional table, supporting up to 100 records per call. (批量获取多维表格数据表中的记录信息,单次调用最多支持查询 100 条记录) + +parameters: + - name: app_token + type: string + required: true + label: + en_US: app_token + zh_Hans: app_token + human_description: + en_US: Unique identifier for the multidimensional table, supports inputting document URL. + zh_Hans: 多维表格的唯一标识符,支持输入文档 URL。 + llm_description: 多维表格的唯一标识符,支持输入文档 URL。 + form: llm + + - name: table_id + type: string + required: false + label: + en_US: table_id + zh_Hans: table_id + human_description: + en_US: Unique identifier for the multidimensional table data, either table_id or table_name must be provided, cannot be empty simultaneously. + zh_Hans: 多维表格数据表的唯一标识符,table_id 和 table_name 至少需要提供一个,不能同时为空。 + llm_description: 多维表格数据表的唯一标识符,table_id 和 table_name 至少需要提供一个,不能同时为空。 + form: llm + + - name: table_name + type: string + required: false + label: + en_US: table_name + zh_Hans: table_name + human_description: + en_US: Name of the multidimensional table data, either table_name or table_id must be provided, cannot be empty simultaneously. + zh_Hans: 多维表格数据表的名称,table_name 和 table_id 至少需要提供一个,不能同时为空。 + llm_description: 多维表格数据表的名称,table_name 和 table_id 至少需要提供一个,不能同时为空。 + form: llm + + - name: record_ids + type: string + required: true + label: + en_US: record_ids + zh_Hans: 记录 ID 列表 + human_description: + en_US: List of record IDs, which can be obtained by calling the "Query Records API". + zh_Hans: 记录 ID 列表,可以通过调用"查询记录接口"获取。 + llm_description: 记录 ID 列表,可以通过调用"查询记录接口"获取。 + form: llm + + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form diff --git a/api/core/tools/provider/builtin/feishu_base/tools/search_records.py b/api/core/tools/provider/builtin/feishu_base/tools/search_records.py new file mode 100644 index 0000000000000000000000000000000000000000..d58b42b82029ced340ccc24ce439a7cc1b7a7354 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/tools/search_records.py @@ -0,0 +1,43 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class SearchRecordsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + if not self.runtime or not self.runtime.credentials: + raise ValueError("Runtime is not set") + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + if not app_id or not app_secret: + raise ValueError("app_id and app_secret are required") + client = FeishuRequest(app_id, app_secret) + + app_token = tool_parameters.get("app_token", "") + table_id = tool_parameters.get("table_id", "") + table_name = tool_parameters.get("table_name", "") + view_id = tool_parameters.get("view_id", "") + field_names = tool_parameters.get("field_names", "") + sort = tool_parameters.get("sort", "") + filters = tool_parameters.get("filter", "") + page_token = tool_parameters.get("page_token", "") + automatic_fields = tool_parameters.get("automatic_fields", False) + user_id_type = tool_parameters.get("user_id_type", "open_id") + page_size = tool_parameters.get("page_size", 20) + + res = client.search_record( + app_token, + table_id, + table_name, + view_id, + field_names, + sort, + filters, + page_token, + automatic_fields, + user_id_type, + page_size, + ) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_base/tools/search_records.yaml b/api/core/tools/provider/builtin/feishu_base/tools/search_records.yaml new file mode 100644 index 0000000000000000000000000000000000000000..decf76d53ed928bc01e87a2ffd922c2e031803af --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/tools/search_records.yaml @@ -0,0 +1,163 @@ +identity: + name: search_records + author: Doug Lea + label: + en_US: Search Records + zh_Hans: 查询记录 +description: + human: + en_US: Query records in a multidimensional table, up to 500 rows per query. + zh_Hans: 查询多维表格数据表中的记录,单次最多查询 500 行记录。 + llm: A tool for querying records in a multidimensional table, up to 500 rows per query. (查询多维表格数据表中的记录,单次最多查询 500 行记录) +parameters: + - name: app_token + type: string + required: true + label: + en_US: app_token + zh_Hans: app_token + human_description: + en_US: Unique identifier for the multidimensional table, supports inputting document URL. + zh_Hans: 多维表格的唯一标识符,支持输入文档 URL。 + llm_description: 多维表格的唯一标识符,支持输入文档 URL。 + form: llm + + - name: table_id + type: string + required: false + label: + en_US: table_id + zh_Hans: table_id + human_description: + en_US: Unique identifier for the multidimensional table data, either table_id or table_name must be provided, cannot be empty simultaneously. + zh_Hans: 多维表格数据表的唯一标识符,table_id 和 table_name 至少需要提供一个,不能同时为空。 + llm_description: 多维表格数据表的唯一标识符,table_id 和 table_name 至少需要提供一个,不能同时为空。 + form: llm + + - name: table_name + type: string + required: false + label: + en_US: table_name + zh_Hans: table_name + human_description: + en_US: Name of the multidimensional table data, either table_name or table_id must be provided, cannot be empty simultaneously. + zh_Hans: 多维表格数据表的名称,table_name 和 table_id 至少需要提供一个,不能同时为空。 + llm_description: 多维表格数据表的名称,table_name 和 table_id 至少需要提供一个,不能同时为空。 + form: llm + + - name: view_id + type: string + required: false + label: + en_US: view_id + zh_Hans: 视图唯一标识 + human_description: + en_US: | + Unique identifier for a view in a multidimensional table. It can be found in the URL's query parameter with the key 'view'. For example: https://svi136aogf123.feishu.cn/base/KWC8bYsYXahYqGsTtqectNn9n3e?table=tblE8a2fmBIEflaE&view=vewlkAVpRx. + zh_Hans: 多维表格中视图的唯一标识,可在多维表格的 URL 地址栏中找到,query 参数中 key 为 view 的部分。例如:https://svi136aogf123.feishu.cn/base/KWC8bYsYXahYqGsTtqectNn9n3e?table=tblE8a2fmBIEflaE&view=vewlkAVpRx。 + llm_description: 多维表格中视图的唯一标识,可在多维表格的 URL 地址栏中找到,query 参数中 key 为 view 的部分。例如:https://svi136aogf123.feishu.cn/base/KWC8bYsYXahYqGsTtqectNn9n3e?table=tblE8a2fmBIEflaE&view=vewlkAVpRx。 + form: llm + + - name: field_names + type: string + required: false + label: + en_US: field_names + zh_Hans: 字段名称 + human_description: + en_US: | + Field names to specify which fields to include in the returned records. Example value: ["Field1", "Field2"]. + zh_Hans: 字段名称,用于指定本次查询返回记录中包含的字段。示例值:["字段1","字段2"]。 + llm_description: 字段名称,用于指定本次查询返回记录中包含的字段。示例值:["字段1","字段2"]。 + form: llm + + - name: sort + type: string + required: false + label: + en_US: sort + zh_Hans: 排序条件 + human_description: + en_US: | + Sorting conditions, for example: [{"field_name":"Multiline Text","desc":true}]. + zh_Hans: 排序条件,例如:[{"field_name":"多行文本","desc":true}]。 + llm_description: 排序条件,例如:[{"field_name":"多行文本","desc":true}]。 + form: llm + + - name: filter + type: string + required: false + label: + en_US: filter + zh_Hans: 筛选条件 + human_description: + en_US: Object containing filter information. For details on how to fill in the filter, refer to the record filter parameter guide (https://open.larkoffice.com/document/uAjLw4CM/ukTMukTMukTM/reference/bitable-v1/app-table-record/record-filter-guide). + zh_Hans: 包含条件筛选信息的对象。了解如何填写 filter,参考记录筛选参数填写指南(https://open.larkoffice.com/document/uAjLw4CM/ukTMukTMukTM/reference/bitable-v1/app-table-record/record-filter-guide)。 + llm_description: 包含条件筛选信息的对象。了解如何填写 filter,参考记录筛选参数填写指南(https://open.larkoffice.com/document/uAjLw4CM/ukTMukTMukTM/reference/bitable-v1/app-table-record/record-filter-guide)。 + form: llm + + - name: automatic_fields + type: boolean + required: false + label: + en_US: automatic_fields + zh_Hans: automatic_fields + human_description: + en_US: Whether to return automatically calculated fields. Default is false, meaning they are not returned. + zh_Hans: 是否返回自动计算的字段。默认为 false,表示不返回。 + llm_description: 是否返回自动计算的字段。默认为 false,表示不返回。 + form: form + + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form + + - name: page_size + type: number + required: false + default: 20 + label: + en_US: page_size + zh_Hans: 分页大小 + human_description: + en_US: | + Page size, default value: 20, maximum value: 500. + zh_Hans: 分页大小,默认值:20,最大值:500。 + llm_description: 分页大小,默认值:20,最大值:500。 + form: form + + - name: page_token + type: string + required: false + label: + en_US: page_token + zh_Hans: 分页标记 + human_description: + en_US: | + Page token, leave empty for the first request to start from the beginning; a new page_token will be returned if there are more items in the paginated query results, which can be used for the next traversal. Example value: "tblsRc9GRRXKqhvW". + zh_Hans: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。示例值:"tblsRc9GRRXKqhvW"。 + llm_description: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。示例值:"tblsRc9GRRXKqhvW"。 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_base/tools/update_records.py b/api/core/tools/provider/builtin/feishu_base/tools/update_records.py new file mode 100644 index 0000000000000000000000000000000000000000..31cf8e18d85b8d7573e050a146eb21bb254ce6e8 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/tools/update_records.py @@ -0,0 +1,25 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class UpdateRecordsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + if not self.runtime or not self.runtime.credentials: + raise ValueError("Runtime is not set") + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + if not app_id or not app_secret: + raise ValueError("app_id and app_secret are required") + client = FeishuRequest(app_id, app_secret) + + app_token = tool_parameters.get("app_token", "") + table_id = tool_parameters.get("table_id", "") + table_name = tool_parameters.get("table_name", "") + records = tool_parameters.get("records", "") + user_id_type = tool_parameters.get("user_id_type", "open_id") + + res = client.update_records(app_token, table_id, table_name, records, user_id_type) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_base/tools/update_records.yaml b/api/core/tools/provider/builtin/feishu_base/tools/update_records.yaml new file mode 100644 index 0000000000000000000000000000000000000000..68117e7136789225bf75724bbdf59c9fbccfbd36 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_base/tools/update_records.yaml @@ -0,0 +1,91 @@ +identity: + name: update_records + author: Doug Lea + label: + en_US: Update Records + zh_Hans: 更新多条记录 +description: + human: + en_US: Update Multiple Records in Multidimensional Table + zh_Hans: 更新多维表格数据表中的多条记录 + llm: A tool for updating multiple records in a multidimensional table. (更新多维表格数据表中的多条记录) +parameters: + - name: app_token + type: string + required: true + label: + en_US: app_token + zh_Hans: app_token + human_description: + en_US: Unique identifier for the multidimensional table, supports inputting document URL. + zh_Hans: 多维表格的唯一标识符,支持输入文档 URL。 + llm_description: 多维表格的唯一标识符,支持输入文档 URL。 + form: llm + + - name: table_id + type: string + required: false + label: + en_US: table_id + zh_Hans: table_id + human_description: + en_US: Unique identifier for the multidimensional table data, either table_id or table_name must be provided, cannot be empty simultaneously. + zh_Hans: 多维表格数据表的唯一标识符,table_id 和 table_name 至少需要提供一个,不能同时为空。 + llm_description: 多维表格数据表的唯一标识符,table_id 和 table_name 至少需要提供一个,不能同时为空。 + form: llm + + - name: table_name + type: string + required: false + label: + en_US: table_name + zh_Hans: table_name + human_description: + en_US: Name of the multidimensional table data, either table_name or table_id must be provided, cannot be empty simultaneously. + zh_Hans: 多维表格数据表的名称,table_name 和 table_id 至少需要提供一个,不能同时为空。 + llm_description: 多维表格数据表的名称,table_name 和 table_id 至少需要提供一个,不能同时为空。 + form: llm + + - name: records + type: string + required: true + label: + en_US: records + zh_Hans: 记录列表 + human_description: + en_US: | + List of records to be updated in this request. Example value: [{"fields":{"multi-line-text":"text content","single_select":"option 1","date":1674206443000},"record_id":"recupK4f4RM5RX"}]. + For supported field types, refer to the integration guide (https://open.larkoffice.com/document/server-docs/docs/bitable-v1/notification). For data structures of different field types, refer to the data structure overview (https://open.larkoffice.com/document/server-docs/docs/bitable-v1/bitable-structure). + zh_Hans: | + 本次请求将要更新的记录列表,示例值:[{"fields":{"多行文本":"文本内容","单选":"选项 1","日期":1674206443000},"record_id":"recupK4f4RM5RX"}]。 + 当前接口支持的字段类型请参考接入指南(https://open.larkoffice.com/document/server-docs/docs/bitable-v1/notification),不同类型字段的数据结构请参考数据结构概述(https://open.larkoffice.com/document/server-docs/docs/bitable-v1/bitable-structure)。 + llm_description: | + 本次请求将要更新的记录列表,示例值:[{"fields":{"多行文本":"文本内容","单选":"选项 1","日期":1674206443000},"record_id":"recupK4f4RM5RX"}]。 + 当前接口支持的字段类型请参考接入指南(https://open.larkoffice.com/document/server-docs/docs/bitable-v1/notification),不同类型字段的数据结构请参考数据结构概述(https://open.larkoffice.com/document/server-docs/docs/bitable-v1/bitable-structure)。 + form: llm + + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form diff --git a/api/core/tools/provider/builtin/feishu_calendar/_assets/icon.png b/api/core/tools/provider/builtin/feishu_calendar/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..2a934747a98c6680065941bcd31d2400da1eaf23 Binary files /dev/null and b/api/core/tools/provider/builtin/feishu_calendar/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/feishu_calendar/feishu_calendar.py b/api/core/tools/provider/builtin/feishu_calendar/feishu_calendar.py new file mode 100644 index 0000000000000000000000000000000000000000..a46a9fa9e80cab7ee8ce061b77f023ba96e6d005 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_calendar/feishu_calendar.py @@ -0,0 +1,7 @@ +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController +from core.tools.utils.feishu_api_utils import auth + + +class FeishuCalendarProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + auth(credentials) diff --git a/api/core/tools/provider/builtin/feishu_calendar/feishu_calendar.yaml b/api/core/tools/provider/builtin/feishu_calendar/feishu_calendar.yaml new file mode 100644 index 0000000000000000000000000000000000000000..db5bab5c1081d99d32e12093a872008cd4251794 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_calendar/feishu_calendar.yaml @@ -0,0 +1,36 @@ +identity: + author: Doug Lea + name: feishu_calendar + label: + en_US: Feishu Calendar + zh_Hans: 飞书日历 + description: + en_US: | + Feishu calendar, requires the following permissions: calendar:calendar:read、calendar:calendar、contact:user.id:readonly. + zh_Hans: | + 飞书日历,需要开通以下权限: calendar:calendar:read、calendar:calendar、contact:user.id:readonly。 + icon: icon.png + tags: + - social + - productivity +credentials_for_provider: + app_id: + type: text-input + required: true + label: + en_US: APP ID + placeholder: + en_US: Please input your feishu app id + zh_Hans: 请输入你的飞书 app id + help: + en_US: Get your app_id and app_secret from Feishu + zh_Hans: 从飞书获取您的 app_id 和 app_secret + url: https://open.larkoffice.com/app + app_secret: + type: secret-input + required: true + label: + en_US: APP Secret + placeholder: + en_US: Please input your app secret + zh_Hans: 请输入你的飞书 app secret diff --git a/api/core/tools/provider/builtin/feishu_calendar/tools/add_event_attendees.py b/api/core/tools/provider/builtin/feishu_calendar/tools/add_event_attendees.py new file mode 100644 index 0000000000000000000000000000000000000000..80287feca176e16b3f433e38493dee87f36f6a4a --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_calendar/tools/add_event_attendees.py @@ -0,0 +1,24 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class AddEventAttendeesTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + if not self.runtime or not self.runtime.credentials: + raise ValueError("Runtime is not set") + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + if not app_id or not app_secret: + raise ValueError("app_id and app_secret are required") + client = FeishuRequest(app_id, app_secret) + + event_id = tool_parameters.get("event_id", "") + attendee_phone_or_email = tool_parameters.get("attendee_phone_or_email", "") + need_notification = tool_parameters.get("need_notification", True) + + res = client.add_event_attendees(event_id, attendee_phone_or_email, need_notification) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_calendar/tools/add_event_attendees.yaml b/api/core/tools/provider/builtin/feishu_calendar/tools/add_event_attendees.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b7744499b073448ba34b57d2edc74f54c5086bdc --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_calendar/tools/add_event_attendees.yaml @@ -0,0 +1,54 @@ +identity: + name: add_event_attendees + author: Doug Lea + label: + en_US: Add Event Attendees + zh_Hans: 添加日程参会人 +description: + human: + en_US: Add Event Attendees + zh_Hans: 添加日程参会人 + llm: A tool for adding attendees to events in Feishu. (在飞书中添加日程参会人) +parameters: + - name: event_id + type: string + required: true + label: + en_US: Event ID + zh_Hans: 日程 ID + human_description: + en_US: | + The ID of the event, which will be returned when the event is created. For example: fb2a6406-26d6-4c8d-a487-6f0246c94d2f_0. + zh_Hans: | + 创建日程时会返回日程 ID。例如: fb2a6406-26d6-4c8d-a487-6f0246c94d2f_0。 + llm_description: | + 日程 ID,创建日程时会返回日程 ID。例如: fb2a6406-26d6-4c8d-a487-6f0246c94d2f_0。 + form: llm + + - name: need_notification + type: boolean + required: false + default: true + label: + en_US: Need Notification + zh_Hans: 是否需要通知 + human_description: + en_US: | + Whether to send a Bot notification to attendees. true: send, false: do not send. + zh_Hans: | + 是否给参与人发送 Bot 通知,true: 发送,false: 不发送。 + llm_description: | + 是否给参与人发送 Bot 通知,true: 发送,false: 不发送。 + form: form + + - name: attendee_phone_or_email + type: string + required: true + label: + en_US: Attendee Phone or Email + zh_Hans: 参会人电话或邮箱 + human_description: + en_US: The list of attendee emails or phone numbers, separated by commas. + zh_Hans: 日程参会人邮箱或者手机号列表,使用逗号分隔。 + llm_description: 日程参会人邮箱或者手机号列表,使用逗号分隔。 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_calendar/tools/create_event.py b/api/core/tools/provider/builtin/feishu_calendar/tools/create_event.py new file mode 100644 index 0000000000000000000000000000000000000000..8820bebdbed922f3b894968d500706725366f6a5 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_calendar/tools/create_event.py @@ -0,0 +1,26 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class CreateEventTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + summary = tool_parameters.get("summary") + description = tool_parameters.get("description") + start_time = tool_parameters.get("start_time") + end_time = tool_parameters.get("end_time") + attendee_ability = tool_parameters.get("attendee_ability") + need_notification = tool_parameters.get("need_notification", True) + auto_record = tool_parameters.get("auto_record", False) + + res = client.create_event( + summary, description, start_time, end_time, attendee_ability, need_notification, auto_record + ) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_calendar/tools/create_event.yaml b/api/core/tools/provider/builtin/feishu_calendar/tools/create_event.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f0784221ce796596ff2c46db9cba7bc14b04e719 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_calendar/tools/create_event.yaml @@ -0,0 +1,119 @@ +identity: + name: create_event + author: Doug Lea + label: + en_US: Create Event + zh_Hans: 创建日程 +description: + human: + en_US: Create Event + zh_Hans: 创建日程 + llm: A tool for creating events in Feishu.(创建飞书日程) +parameters: + - name: summary + type: string + required: false + label: + en_US: Summary + zh_Hans: 日程标题 + human_description: + en_US: The title of the event. If not filled, the event title will display (No Subject). + zh_Hans: 日程标题,若不填则日程标题显示 (无主题)。 + llm_description: 日程标题,若不填则日程标题显示 (无主题)。 + form: llm + + - name: description + type: string + required: false + label: + en_US: Description + zh_Hans: 日程描述 + human_description: + en_US: The description of the event. + zh_Hans: 日程描述。 + llm_description: 日程描述。 + form: llm + + - name: need_notification + type: boolean + required: false + default: true + label: + en_US: Need Notification + zh_Hans: 是否发送通知 + human_description: + en_US: | + Whether to send a bot message when the event is created, true: send, false: do not send. + zh_Hans: 创建日程时是否发送 bot 消息,true:发送,false:不发送。 + llm_description: 创建日程时是否发送 bot 消息,true:发送,false:不发送。 + form: form + + - name: start_time + type: string + required: true + label: + en_US: Start Time + zh_Hans: 开始时间 + human_description: + en_US: | + The start time of the event, format: 2006-01-02 15:04:05. + zh_Hans: 日程开始时间,格式:2006-01-02 15:04:05。 + llm_description: 日程开始时间,格式:2006-01-02 15:04:05。 + form: llm + + - name: end_time + type: string + required: true + label: + en_US: End Time + zh_Hans: 结束时间 + human_description: + en_US: | + The end time of the event, format: 2006-01-02 15:04:05. + zh_Hans: 日程结束时间,格式:2006-01-02 15:04:05。 + llm_description: 日程结束时间,格式:2006-01-02 15:04:05。 + form: llm + + - name: attendee_ability + type: select + required: false + options: + - value: none + label: + en_US: none + zh_Hans: 无 + - value: can_see_others + label: + en_US: can_see_others + zh_Hans: 可以查看参与人列表 + - value: can_invite_others + label: + en_US: can_invite_others + zh_Hans: 可以邀请其它参与人 + - value: can_modify_event + label: + en_US: can_modify_event + zh_Hans: 可以编辑日程 + default: "none" + label: + en_US: attendee_ability + zh_Hans: 参会人权限 + human_description: + en_US: Attendee ability, optional values are none, can_see_others, can_invite_others, can_modify_event, with a default value of none. + zh_Hans: 参会人权限,可选值有无、可以查看参与人列表、可以邀请其它参与人、可以编辑日程,默认值为无。 + llm_description: 参会人权限,可选值有无、可以查看参与人列表、可以邀请其它参与人、可以编辑日程,默认值为无。 + form: form + + - name: auto_record + type: boolean + required: false + default: false + label: + en_US: Auto Record + zh_Hans: 自动录制 + human_description: + en_US: | + Whether to enable automatic recording, true: enabled, automatically record when the meeting starts; false: not enabled. + zh_Hans: 是否开启自动录制,true:开启,会议开始后自动录制;false:不开启。 + llm_description: 是否开启自动录制,true:开启,会议开始后自动录制;false:不开启。 + form: form diff --git a/api/core/tools/provider/builtin/feishu_calendar/tools/delete_event.py b/api/core/tools/provider/builtin/feishu_calendar/tools/delete_event.py new file mode 100644 index 0000000000000000000000000000000000000000..02e9b445219ac8f622568672e327ad52228d2c79 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_calendar/tools/delete_event.py @@ -0,0 +1,23 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class DeleteEventTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + if not self.runtime or not self.runtime.credentials: + raise ValueError("Runtime is not set") + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + if not app_id or not app_secret: + raise ValueError("app_id and app_secret are required") + client = FeishuRequest(app_id, app_secret) + + event_id = tool_parameters.get("event_id", "") + need_notification = tool_parameters.get("need_notification", True) + + res = client.delete_event(event_id, need_notification) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_calendar/tools/delete_event.yaml b/api/core/tools/provider/builtin/feishu_calendar/tools/delete_event.yaml new file mode 100644 index 0000000000000000000000000000000000000000..54fdb04acc33717caef660e5feedcb8b61ef3360 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_calendar/tools/delete_event.yaml @@ -0,0 +1,38 @@ +identity: + name: delete_event + author: Doug Lea + label: + en_US: Delete Event + zh_Hans: 删除日程 +description: + human: + en_US: Delete Event + zh_Hans: 删除日程 + llm: A tool for deleting events in Feishu.(在飞书中删除日程) +parameters: + - name: event_id + type: string + required: true + label: + en_US: Event ID + zh_Hans: 日程 ID + human_description: + en_US: | + The ID of the event, for example: e8b9791c-39ae-4908-8ad8-66b13159b9fb_0. + zh_Hans: 日程 ID,例如:e8b9791c-39ae-4908-8ad8-66b13159b9fb_0。 + llm_description: 日程 ID,例如:e8b9791c-39ae-4908-8ad8-66b13159b9fb_0。 + form: llm + + - name: need_notification + type: boolean + required: false + default: true + label: + en_US: Need Notification + zh_Hans: 是否需要通知 + human_description: + en_US: | + Indicates whether to send bot notifications to event participants upon deletion. true: send, false: do not send. + zh_Hans: 删除日程是否给日程参与人发送 bot 通知,true:发送,false:不发送。 + llm_description: 删除日程是否给日程参与人发送 bot 通知,true:发送,false:不发送。 + form: form diff --git a/api/core/tools/provider/builtin/feishu_calendar/tools/get_primary_calendar.py b/api/core/tools/provider/builtin/feishu_calendar/tools/get_primary_calendar.py new file mode 100644 index 0000000000000000000000000000000000000000..4dafe4b3baf0cdc90a15291d314ddc2944252cee --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_calendar/tools/get_primary_calendar.py @@ -0,0 +1,22 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class GetPrimaryCalendarTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + if not self.runtime or not self.runtime.credentials: + raise ValueError("Runtime is not set") + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + if not app_id or not app_secret: + raise ValueError("app_id and app_secret are required") + client = FeishuRequest(app_id, app_secret) + + user_id_type = tool_parameters.get("user_id_type", "open_id") + + res = client.get_primary_calendar(user_id_type) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_calendar/tools/get_primary_calendar.yaml b/api/core/tools/provider/builtin/feishu_calendar/tools/get_primary_calendar.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3440c85d4a97334b47e2b145af4fc6edc824fb3a --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_calendar/tools/get_primary_calendar.yaml @@ -0,0 +1,37 @@ +identity: + name: get_primary_calendar + author: Doug Lea + label: + en_US: Get Primary Calendar + zh_Hans: 查询主日历信息 +description: + human: + en_US: Get Primary Calendar + zh_Hans: 查询主日历信息 + llm: A tool for querying primary calendar information in Feishu.(在飞书中查询主日历信息) +parameters: + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form diff --git a/api/core/tools/provider/builtin/feishu_calendar/tools/list_events.py b/api/core/tools/provider/builtin/feishu_calendar/tools/list_events.py new file mode 100644 index 0000000000000000000000000000000000000000..2e8ca968b3cc423696ffdbc320903b6ee281c247 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_calendar/tools/list_events.py @@ -0,0 +1,25 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class ListEventsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + if not self.runtime or not self.runtime.credentials: + raise ValueError("Runtime is not set") + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + if not app_id or not app_secret: + raise ValueError("app_id and app_secret are required") + client = FeishuRequest(app_id, app_secret) + + start_time = tool_parameters.get("start_time", "") + end_time = tool_parameters.get("end_time", "") + page_token = tool_parameters.get("page_token", "") + page_size = tool_parameters.get("page_size", 50) + + res = client.list_events(start_time, end_time, page_token, page_size) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_calendar/tools/list_events.yaml b/api/core/tools/provider/builtin/feishu_calendar/tools/list_events.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5f0155a24658661379a04d55b90687cebd73a1e7 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_calendar/tools/list_events.yaml @@ -0,0 +1,62 @@ +identity: + name: list_events + author: Doug Lea + label: + en_US: List Events + zh_Hans: 获取日程列表 +description: + human: + en_US: List Events + zh_Hans: 获取日程列表 + llm: A tool for listing events in Feishu.(在飞书中获取日程列表) +parameters: + - name: start_time + type: string + required: false + label: + en_US: Start Time + zh_Hans: 开始时间 + human_description: + en_US: | + The start time, defaults to 0:00 of the current day if not provided, format: 2006-01-02 15:04:05. + zh_Hans: 开始时间,不传值时默认当天 0 点时间,格式为:2006-01-02 15:04:05。 + llm_description: 开始时间,不传值时默认当天 0 点时间,格式为:2006-01-02 15:04:05。 + form: llm + + - name: end_time + type: string + required: false + label: + en_US: End Time + zh_Hans: 结束时间 + human_description: + en_US: | + The end time, defaults to 23:59 of the current day if not provided, format: 2006-01-02 15:04:05. + zh_Hans: 结束时间,不传值时默认当天 23:59 分时间,格式为:2006-01-02 15:04:05。 + llm_description: 结束时间,不传值时默认当天 23:59 分时间,格式为:2006-01-02 15:04:05。 + form: llm + + - name: page_size + type: number + required: false + default: 50 + label: + en_US: Page Size + zh_Hans: 分页大小 + human_description: + en_US: The page size, i.e., the number of data entries returned in a single request. The default value is 50, and the value range is [50,1000]. + zh_Hans: 分页大小,即单次请求所返回的数据条目数。默认值为 50,取值范围为 [50,1000]。 + llm_description: 分页大小,即单次请求所返回的数据条目数。默认值为 50,取值范围为 [50,1000]。 + form: form + + - name: page_token + type: string + required: false + label: + en_US: Page Token + zh_Hans: 分页标记 + human_description: + en_US: The pagination token. Leave it blank for the first request, indicating to start traversing from the beginning; when the pagination query result has more items, a new page_token will be returned simultaneously, which can be used to obtain the query result in the next traversal. + zh_Hans: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + llm_description: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_calendar/tools/search_events.py b/api/core/tools/provider/builtin/feishu_calendar/tools/search_events.py new file mode 100644 index 0000000000000000000000000000000000000000..dc365205a4cffa34e25c59a66739afc23471d21a --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_calendar/tools/search_events.py @@ -0,0 +1,23 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class SearchEventsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + query = tool_parameters.get("query") + start_time = tool_parameters.get("start_time") + end_time = tool_parameters.get("end_time") + page_token = tool_parameters.get("page_token") + user_id_type = tool_parameters.get("user_id_type", "open_id") + page_size = tool_parameters.get("page_size", 20) + + res = client.search_events(query, start_time, end_time, page_token, user_id_type, page_size) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_calendar/tools/search_events.yaml b/api/core/tools/provider/builtin/feishu_calendar/tools/search_events.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bd60a07b5b534143aa69f25a5bf122e13e591e9c --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_calendar/tools/search_events.yaml @@ -0,0 +1,100 @@ +identity: + name: search_events + author: Doug Lea + label: + en_US: Search Events + zh_Hans: 搜索日程 +description: + human: + en_US: Search Events + zh_Hans: 搜索日程 + llm: A tool for searching events in Feishu.(在飞书中搜索日程) +parameters: + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form + + - name: query + type: string + required: true + label: + en_US: Query + zh_Hans: 搜索关键字 + human_description: + en_US: The search keyword used for fuzzy searching event names, with a maximum input of 200 characters. + zh_Hans: 用于模糊查询日程名称的搜索关键字,最大输入 200 字符。 + llm_description: 用于模糊查询日程名称的搜索关键字,最大输入 200 字符。 + form: llm + + - name: start_time + type: string + required: false + label: + en_US: Start Time + zh_Hans: 开始时间 + human_description: + en_US: | + The start time, defaults to 0:00 of the current day if not provided, format: 2006-01-02 15:04:05. + zh_Hans: 开始时间,不传值时默认当天 0 点时间,格式为:2006-01-02 15:04:05。 + llm_description: 开始时间,不传值时默认当天 0 点时间,格式为:2006-01-02 15:04:05。 + form: llm + + - name: end_time + type: string + required: false + label: + en_US: End Time + zh_Hans: 结束时间 + human_description: + en_US: | + The end time, defaults to 23:59 of the current day if not provided, format: 2006-01-02 15:04:05. + zh_Hans: 结束时间,不传值时默认当天 23:59 分时间,格式为:2006-01-02 15:04:05。 + llm_description: 结束时间,不传值时默认当天 23:59 分时间,格式为:2006-01-02 15:04:05。 + form: llm + + - name: page_size + type: number + required: false + default: 20 + label: + en_US: Page Size + zh_Hans: 分页大小 + human_description: + en_US: The page size, i.e., the number of data entries returned in a single request. The default value is 20, and the value range is [10,100]. + zh_Hans: 分页大小,即单次请求所返回的数据条目数。默认值为 20,取值范围为 [10,100]。 + llm_description: 分页大小,即单次请求所返回的数据条目数。默认值为 20,取值范围为 [10,100]。 + form: form + + - name: page_token + type: string + required: false + label: + en_US: Page Token + zh_Hans: 分页标记 + human_description: + en_US: The pagination token. Leave it blank for the first request, indicating to start traversing from the beginning; when the pagination query result has more items, a new page_token will be returned simultaneously, which can be used to obtain the query result in the next traversal. + zh_Hans: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + llm_description: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_calendar/tools/update_event.py b/api/core/tools/provider/builtin/feishu_calendar/tools/update_event.py new file mode 100644 index 0000000000000000000000000000000000000000..b20eb6c31828e45be76c83ebf9181cfbd7234a77 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_calendar/tools/update_event.py @@ -0,0 +1,28 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class UpdateEventTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + if not self.runtime or not self.runtime.credentials: + raise ValueError("Runtime is not set") + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + if not app_id or not app_secret: + raise ValueError("app_id and app_secret are required") + client = FeishuRequest(app_id, app_secret) + + event_id = tool_parameters.get("event_id", "") + summary = tool_parameters.get("summary", "") + description = tool_parameters.get("description", "") + need_notification = tool_parameters.get("need_notification", True) + start_time = tool_parameters.get("start_time", "") + end_time = tool_parameters.get("end_time", "") + auto_record = tool_parameters.get("auto_record", False) + + res = client.update_event(event_id, summary, description, need_notification, start_time, end_time, auto_record) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_calendar/tools/update_event.yaml b/api/core/tools/provider/builtin/feishu_calendar/tools/update_event.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4d60dbf8c8e1b0cbf7022decba9421c00ac6095f --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_calendar/tools/update_event.yaml @@ -0,0 +1,100 @@ +identity: + name: update_event + author: Doug Lea + label: + en_US: Update Event + zh_Hans: 更新日程 +description: + human: + en_US: Update Event + zh_Hans: 更新日程 + llm: A tool for updating events in Feishu.(更新飞书中的日程) +parameters: + - name: event_id + type: string + required: true + label: + en_US: Event ID + zh_Hans: 日程 ID + human_description: + en_US: | + The ID of the event, for example: e8b9791c-39ae-4908-8ad8-66b13159b9fb_0. + zh_Hans: 日程 ID,例如:e8b9791c-39ae-4908-8ad8-66b13159b9fb_0。 + llm_description: 日程 ID,例如:e8b9791c-39ae-4908-8ad8-66b13159b9fb_0。 + form: llm + + - name: summary + type: string + required: false + label: + en_US: Summary + zh_Hans: 日程标题 + human_description: + en_US: The title of the event. + zh_Hans: 日程标题。 + llm_description: 日程标题。 + form: llm + + - name: description + type: string + required: false + label: + en_US: Description + zh_Hans: 日程描述 + human_description: + en_US: The description of the event. + zh_Hans: 日程描述。 + llm_description: 日程描述。 + form: llm + + - name: need_notification + type: boolean + required: false + label: + en_US: Need Notification + zh_Hans: 是否发送通知 + human_description: + en_US: | + Whether to send a bot message when the event is updated, true: send, false: do not send. + zh_Hans: 更新日程时是否发送 bot 消息,true:发送,false:不发送。 + llm_description: 更新日程时是否发送 bot 消息,true:发送,false:不发送。 + form: form + + - name: start_time + type: string + required: false + label: + en_US: Start Time + zh_Hans: 开始时间 + human_description: + en_US: | + The start time of the event, format: 2006-01-02 15:04:05. + zh_Hans: 日程开始时间,格式:2006-01-02 15:04:05。 + llm_description: 日程开始时间,格式:2006-01-02 15:04:05。 + form: llm + + - name: end_time + type: string + required: false + label: + en_US: End Time + zh_Hans: 结束时间 + human_description: + en_US: | + The end time of the event, format: 2006-01-02 15:04:05. + zh_Hans: 日程结束时间,格式:2006-01-02 15:04:05。 + llm_description: 日程结束时间,格式:2006-01-02 15:04:05。 + form: llm + + - name: auto_record + type: boolean + required: false + label: + en_US: Auto Record + zh_Hans: 自动录制 + human_description: + en_US: | + Whether to enable automatic recording, true: enabled, automatically record when the meeting starts; false: not enabled. + zh_Hans: 是否开启自动录制,true:开启,会议开始后自动录制;false:不开启。 + llm_description: 是否开启自动录制,true:开启,会议开始后自动录制;false:不开启。 + form: form diff --git a/api/core/tools/provider/builtin/feishu_document/_assets/icon.svg b/api/core/tools/provider/builtin/feishu_document/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..5a0a6416b3db3205b2e8c5d7039af120cfdd5b07 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_document/_assets/icon.svg @@ -0,0 +1,9 @@ + + + + + + + + + diff --git a/api/core/tools/provider/builtin/feishu_document/feishu_document.py b/api/core/tools/provider/builtin/feishu_document/feishu_document.py new file mode 100644 index 0000000000000000000000000000000000000000..217ae52082b82cfc0eb60b969de17a9552b42da8 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_document/feishu_document.py @@ -0,0 +1,7 @@ +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController +from core.tools.utils.feishu_api_utils import auth + + +class FeishuDocumentProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + auth(credentials) diff --git a/api/core/tools/provider/builtin/feishu_document/feishu_document.yaml b/api/core/tools/provider/builtin/feishu_document/feishu_document.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8f9afa6149445c66ada6171f7c5cd1465a080fb1 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_document/feishu_document.yaml @@ -0,0 +1,36 @@ +identity: + author: Doug Lea + name: feishu_document + label: + en_US: Lark Cloud Document + zh_Hans: 飞书云文档 + description: + en_US: | + Lark cloud document, requires the following permissions: docx:document、drive:drive、docs:document.content:read. + zh_Hans: | + 飞书云文档,需要开通以下权限: docx:document、drive:drive、docs:document.content:read。 + icon: icon.svg + tags: + - social + - productivity +credentials_for_provider: + app_id: + type: text-input + required: true + label: + en_US: APP ID + placeholder: + en_US: Please input your feishu app id + zh_Hans: 请输入你的飞书 app id + help: + en_US: Get your app_id and app_secret from Feishu + zh_Hans: 从飞书获取您的 app_id 和 app_secret + url: https://open.larkoffice.com/app + app_secret: + type: secret-input + required: true + label: + en_US: APP Secret + placeholder: + en_US: Please input your app secret + zh_Hans: 请输入你的飞书 app secret diff --git a/api/core/tools/provider/builtin/feishu_document/tools/create_document.py b/api/core/tools/provider/builtin/feishu_document/tools/create_document.py new file mode 100644 index 0000000000000000000000000000000000000000..1533f594172878f46311889ac447d7d0bc8c4d2d --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_document/tools/create_document.py @@ -0,0 +1,23 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class CreateDocumentTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + if not self.runtime or not self.runtime.credentials: + raise ValueError("Runtime is not set") + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + if not app_id or not app_secret: + raise ValueError("app_id and app_secret are required") + client = FeishuRequest(app_id, app_secret) + + title = tool_parameters.get("title", "") + content = tool_parameters.get("content", "") + folder_token = tool_parameters.get("folder_token", "") + + res = client.create_document(title, content, folder_token) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_document/tools/create_document.yaml b/api/core/tools/provider/builtin/feishu_document/tools/create_document.yaml new file mode 100644 index 0000000000000000000000000000000000000000..85382e9d8e8d1f4bd3f72eb7cb62c2584a0c4f6d --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_document/tools/create_document.yaml @@ -0,0 +1,48 @@ +identity: + name: create_document + author: Doug Lea + label: + en_US: Create Lark document + zh_Hans: 创建飞书文档 +description: + human: + en_US: Create Lark document + zh_Hans: 创建飞书文档,支持创建空文档和带内容的文档,支持 markdown 语法创建。应用需要开启机器人能力(https://open.feishu.cn/document/faq/trouble-shooting/how-to-enable-bot-ability)。 + llm: A tool for creating Feishu documents. +parameters: + - name: title + type: string + required: false + label: + en_US: Document title + zh_Hans: 文档标题 + human_description: + en_US: Document title, only supports plain text content. + zh_Hans: 文档标题,只支持纯文本内容。 + llm_description: 文档标题,只支持纯文本内容,可以为空。 + form: llm + + - name: content + type: string + required: false + label: + en_US: Document content + zh_Hans: 文档内容 + human_description: + en_US: Document content, supports markdown syntax, can be empty. + zh_Hans: 文档内容,支持 markdown 语法,可以为空。 + llm_description: 文档内容,支持 markdown 语法,可以为空。 + form: llm + + - name: folder_token + type: string + required: false + label: + en_US: folder_token + zh_Hans: 文档所在文件夹的 Token + human_description: + en_US: | + The token of the folder where the document is located. If it is not passed or is empty, it means the root directory. For Example: https://svi136aogf123.feishu.cn/drive/folder/JgR9fiG9AlPt8EdsSNpcGjIInbf + zh_Hans: 文档所在文件夹的 Token,不传或传空表示根目录。例如:https://svi136aogf123.feishu.cn/drive/folder/JgR9fiG9AlPt8EdsSNpcGjIInbf。 + llm_description: 文档所在文件夹的 Token,不传或传空表示根目录。例如:https://svi136aogf123.feishu.cn/drive/folder/JgR9fiG9AlPt8EdsSNpcGjIInbf。 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_document/tools/get_document_content.py b/api/core/tools/provider/builtin/feishu_document/tools/get_document_content.py new file mode 100644 index 0000000000000000000000000000000000000000..e67a017facc8d47f517cace5e265ab52b152bb0e --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_document/tools/get_document_content.py @@ -0,0 +1,19 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class GetDocumentRawContentTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + document_id = tool_parameters.get("document_id") + mode = tool_parameters.get("mode", "markdown") + lang = tool_parameters.get("lang", "0") + + res = client.get_document_content(document_id, mode, lang) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_document/tools/get_document_content.yaml b/api/core/tools/provider/builtin/feishu_document/tools/get_document_content.yaml new file mode 100644 index 0000000000000000000000000000000000000000..15e827cde91ee69497e047684bbc61888a681385 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_document/tools/get_document_content.yaml @@ -0,0 +1,70 @@ +identity: + name: get_document_content + author: Doug Lea + label: + en_US: Get Document Content + zh_Hans: 获取飞书云文档的内容 +description: + human: + en_US: Get document content + zh_Hans: 获取飞书云文档的内容 + llm: A tool for retrieving content from Feishu cloud documents. +parameters: + - name: document_id + type: string + required: true + label: + en_US: document_id + zh_Hans: 飞书文档的唯一标识 + human_description: + en_US: Unique identifier for a Feishu document. You can also input the document's URL. + zh_Hans: 飞书文档的唯一标识,支持输入文档的 URL。 + llm_description: 飞书文档的唯一标识,支持输入文档的 URL。 + form: llm + + - name: mode + type: select + required: false + options: + - value: text + label: + en_US: text + zh_Hans: text + - value: markdown + label: + en_US: markdown + zh_Hans: markdown + default: "markdown" + label: + en_US: mode + zh_Hans: 文档返回格式 + human_description: + en_US: Format of the document return, optional values are text, markdown, can be empty, default is markdown. + zh_Hans: 文档返回格式,可选值有 text、markdown,可以为空,默认值为 markdown。 + llm_description: 文档返回格式,可选值有 text、markdown,可以为空,默认值为 markdown。 + form: form + + - name: lang + type: select + required: false + options: + - value: "0" + label: + en_US: User's default name + zh_Hans: 用户的默认名称 + - value: "1" + label: + en_US: User's English name + zh_Hans: 用户的英文名称 + default: "0" + label: + en_US: lang + zh_Hans: 指定@用户的语言 + human_description: + en_US: | + Specifies the language for MentionUser, optional values are [0, 1]. 0: User's default name, 1: User's English name, default is 0. + zh_Hans: | + 指定返回的 MentionUser,即@用户的语言,可选值有 [0,1]。0: 该用户的默认名称,1: 该用户的英文名称,默认值为 0。 + llm_description: | + 指定返回的 MentionUser,即@用户的语言,可选值有 [0,1]。0: 该用户的默认名称,1: 该用户的英文名称,默认值为 0。 + form: form diff --git a/api/core/tools/provider/builtin/feishu_document/tools/list_document_blocks.py b/api/core/tools/provider/builtin/feishu_document/tools/list_document_blocks.py new file mode 100644 index 0000000000000000000000000000000000000000..8ea68a2ed878552feef82525fdc06c573f254cd9 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_document/tools/list_document_blocks.py @@ -0,0 +1,24 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class ListDocumentBlockTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + if not self.runtime or not self.runtime.credentials: + raise ValueError("Runtime is not set") + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + if not app_id or not app_secret: + raise ValueError("app_id and app_secret are required") + client = FeishuRequest(app_id, app_secret) + + document_id = tool_parameters.get("document_id", "") + page_token = tool_parameters.get("page_token", "") + user_id_type = tool_parameters.get("user_id_type", "open_id") + page_size = tool_parameters.get("page_size", 500) + + res = client.list_document_blocks(document_id, page_token, user_id_type, page_size) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_document/tools/list_document_blocks.yaml b/api/core/tools/provider/builtin/feishu_document/tools/list_document_blocks.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d4fab96c1f9601d602838f028f96050d131fc5cb --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_document/tools/list_document_blocks.yaml @@ -0,0 +1,74 @@ +identity: + name: list_document_blocks + author: Doug Lea + label: + en_US: List Document Blocks + zh_Hans: 获取飞书文档所有块 +description: + human: + en_US: List document blocks + zh_Hans: 获取飞书文档所有块的富文本内容并分页返回 + llm: A tool to get all blocks of Feishu documents +parameters: + - name: document_id + type: string + required: true + label: + en_US: document_id + zh_Hans: 飞书文档的唯一标识 + human_description: + en_US: Unique identifier for a Feishu document. You can also input the document's URL. + zh_Hans: 飞书文档的唯一标识,支持输入文档的 URL。 + llm_description: 飞书文档的唯一标识,支持输入文档的 URL。 + form: llm + + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form + + - name: page_size + type: number + required: false + default: 500 + label: + en_US: page_size + zh_Hans: 分页大小 + human_description: + en_US: Paging size, the default and maximum value is 500. + zh_Hans: 分页大小, 默认值和最大值为 500。 + llm_description: 分页大小, 表示一次请求最多返回多少条数据,默认值和最大值为 500。 + form: form + + - name: page_token + type: string + required: false + label: + en_US: page_token + zh_Hans: 分页标记 + human_description: + en_US: Pagination token used to navigate through query results, allowing retrieval of additional items in subsequent requests. + zh_Hans: 分页标记,用于分页查询结果,以便下次遍历时获取更多项。 + llm_description: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_document/tools/write_document.py b/api/core/tools/provider/builtin/feishu_document/tools/write_document.py new file mode 100644 index 0000000000000000000000000000000000000000..59f08f53dc68de649099616e04c9235cddfaa354 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_document/tools/write_document.py @@ -0,0 +1,19 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class CreateDocumentTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + document_id = tool_parameters.get("document_id") + content = tool_parameters.get("content") + position = tool_parameters.get("position", "end") + + res = client.write_document(document_id, content, position) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_document/tools/write_document.yaml b/api/core/tools/provider/builtin/feishu_document/tools/write_document.yaml new file mode 100644 index 0000000000000000000000000000000000000000..de70f4e7726a28a797663b166d8ed21ab1322777 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_document/tools/write_document.yaml @@ -0,0 +1,57 @@ +identity: + name: write_document + author: Doug Lea + label: + en_US: Write Document + zh_Hans: 在飞书文档中新增内容 +description: + human: + en_US: Adding new content to Lark documents + zh_Hans: 在飞书文档中新增内容 + llm: A tool for adding new content to Lark documents. +parameters: + - name: document_id + type: string + required: true + label: + en_US: document_id + zh_Hans: 飞书文档的唯一标识 + human_description: + en_US: Unique identifier for a Feishu document. You can also input the document's URL. + zh_Hans: 飞书文档的唯一标识,支持输入文档的 URL。 + llm_description: 飞书文档的唯一标识,支持输入文档的 URL。 + form: llm + + - name: content + type: string + required: true + label: + en_US: Plain text or Markdown content + zh_Hans: 纯文本或 Markdown 内容 + human_description: + en_US: Plain text or Markdown content. Note that embedded tables in the document should not have merged cells. + zh_Hans: 纯文本或 Markdown 内容。注意文档的内嵌套表格不允许有单元格合并。 + llm_description: 纯文本或 Markdown 内容,注意文档的内嵌套表格不允许有单元格合并。 + form: llm + + - name: position + type: select + required: false + options: + - value: start + label: + en_US: document start + zh_Hans: 文档开始 + - value: end + label: + en_US: document end + zh_Hans: 文档结束 + default: "end" + label: + en_US: position + zh_Hans: 内容添加位置 + human_description: + en_US: Content insertion position, optional values are start, end. 'start' means adding content at the beginning of the document; 'end' means adding content at the end of the document. The default value is end. + zh_Hans: 内容添加位置,可选值有 start、end。start 表示在文档开头添加内容;end 表示在文档结尾添加内容,默认值为 end。 + llm_description: 内容添加位置,可选值有 start、end。start 表示在文档开头添加内容;end 表示在文档结尾添加内容,默认值为 end。 + form: form diff --git a/api/core/tools/provider/builtin/feishu_message/_assets/icon.svg b/api/core/tools/provider/builtin/feishu_message/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..222a1571f9bbbbb48a3dc7450386bb056939804a --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_message/_assets/icon.svg @@ -0,0 +1,19 @@ + + + + diff --git a/api/core/tools/provider/builtin/feishu_message/feishu_message.py b/api/core/tools/provider/builtin/feishu_message/feishu_message.py new file mode 100644 index 0000000000000000000000000000000000000000..a3b54737691c9cec454b239a7feedfee4327f282 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_message/feishu_message.py @@ -0,0 +1,7 @@ +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController +from core.tools.utils.feishu_api_utils import auth + + +class FeishuMessageProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + auth(credentials) diff --git a/api/core/tools/provider/builtin/feishu_message/feishu_message.yaml b/api/core/tools/provider/builtin/feishu_message/feishu_message.yaml new file mode 100644 index 0000000000000000000000000000000000000000..56683ec1680f4026745a37d757f9ce16b83354ac --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_message/feishu_message.yaml @@ -0,0 +1,36 @@ +identity: + author: Doug Lea + name: feishu_message + label: + en_US: Lark Message + zh_Hans: 飞书消息 + description: + en_US: | + Lark message, requires the following permissions: im:message、im:message.group_msg. + zh_Hans: | + 飞书消息,需要开通以下权限: im:message、im:message.group_msg。 + icon: icon.svg + tags: + - social + - productivity +credentials_for_provider: + app_id: + type: text-input + required: true + label: + en_US: APP ID + placeholder: + en_US: Please input your feishu app id + zh_Hans: 请输入你的飞书 app id + help: + en_US: Get your app_id and app_secret from Feishu + zh_Hans: 从飞书获取您的 app_id 和 app_secret + url: https://open.larkoffice.com/app + app_secret: + type: secret-input + required: true + label: + en_US: APP Secret + placeholder: + en_US: Please input your app secret + zh_Hans: 请输入你的飞书 app secret diff --git a/api/core/tools/provider/builtin/feishu_message/tools/get_chat_messages.py b/api/core/tools/provider/builtin/feishu_message/tools/get_chat_messages.py new file mode 100644 index 0000000000000000000000000000000000000000..7eb29230b2ceb020ab55e6f3c0c7af968fc10b3c --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_message/tools/get_chat_messages.py @@ -0,0 +1,23 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class GetChatMessagesTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + container_id = tool_parameters.get("container_id") + start_time = tool_parameters.get("start_time") + end_time = tool_parameters.get("end_time") + page_token = tool_parameters.get("page_token") + sort_type = tool_parameters.get("sort_type", "ByCreateTimeAsc") + page_size = tool_parameters.get("page_size", 20) + + res = client.get_chat_messages(container_id, start_time, end_time, page_token, sort_type, page_size) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_message/tools/get_chat_messages.yaml b/api/core/tools/provider/builtin/feishu_message/tools/get_chat_messages.yaml new file mode 100644 index 0000000000000000000000000000000000000000..984c9120e8cd9686d2a99dc8fb193f9437f78a26 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_message/tools/get_chat_messages.yaml @@ -0,0 +1,96 @@ +identity: + name: get_chat_messages + author: Doug Lea + label: + en_US: Get Chat Messages + zh_Hans: 获取指定单聊、群聊的消息历史 +description: + human: + en_US: Get Chat Messages + zh_Hans: 获取指定单聊、群聊的消息历史 + llm: A tool for getting chat messages from specific one-on-one chats or group chats.(获取指定单聊、群聊的消息历史) +parameters: + - name: container_id + type: string + required: true + label: + en_US: Container Id + zh_Hans: 群聊或单聊的 ID + human_description: + en_US: The ID of the group chat or single chat. Refer to the group ID description for how to obtain it. https://open.feishu.cn/document/server-docs/group/chat/chat-id-description + zh_Hans: 群聊或单聊的 ID,获取方式参见群 ID 说明。https://open.feishu.cn/document/server-docs/group/chat/chat-id-description + llm_description: 群聊或单聊的 ID,获取方式参见群 ID 说明。https://open.feishu.cn/document/server-docs/group/chat/chat-id-description + form: llm + + - name: start_time + type: string + required: false + label: + en_US: Start Time + zh_Hans: 起始时间 + human_description: + en_US: The start time for querying historical messages, formatted as "2006-01-02 15:04:05". + zh_Hans: 待查询历史信息的起始时间,格式为 "2006-01-02 15:04:05"。 + llm_description: 待查询历史信息的起始时间,格式为 "2006-01-02 15:04:05"。 + form: llm + + - name: end_time + type: string + required: false + label: + en_US: End Time + zh_Hans: 结束时间 + human_description: + en_US: The end time for querying historical messages, formatted as "2006-01-02 15:04:05". + zh_Hans: 待查询历史信息的结束时间,格式为 "2006-01-02 15:04:05"。 + llm_description: 待查询历史信息的结束时间,格式为 "2006-01-02 15:04:05"。 + form: llm + + - name: sort_type + type: select + required: false + options: + - value: ByCreateTimeAsc + label: + en_US: ByCreateTimeAsc + zh_Hans: ByCreateTimeAsc + - value: ByCreateTimeDesc + label: + en_US: ByCreateTimeDesc + zh_Hans: ByCreateTimeDesc + default: "ByCreateTimeAsc" + label: + en_US: Sort Type + zh_Hans: 排序方式 + human_description: + en_US: | + The message sorting method. Optional values are ByCreateTimeAsc: sorted in ascending order by message creation time; ByCreateTimeDesc: sorted in descending order by message creation time. The default value is ByCreateTimeAsc. Note: When using page_token for pagination requests, the sorting method (sort_type) is consistent with the first request and cannot be changed midway. + zh_Hans: | + 消息排序方式,可选值有 ByCreateTimeAsc:按消息创建时间升序排列;ByCreateTimeDesc:按消息创建时间降序排列。默认值为:ByCreateTimeAsc。注意:使用 page_token 分页请求时,排序方式(sort_type)均与第一次请求一致,不支持中途改换排序方式。 + llm_description: 消息排序方式,可选值有 ByCreateTimeAsc:按消息创建时间升序排列;ByCreateTimeDesc:按消息创建时间降序排列。默认值为:ByCreateTimeAsc。注意:使用 page_token 分页请求时,排序方式(sort_type)均与第一次请求一致,不支持中途改换排序方式。 + form: form + + - name: page_size + type: number + required: false + default: 20 + label: + en_US: Page Size + zh_Hans: 分页大小 + human_description: + en_US: The page size, i.e., the number of data entries returned in a single request. The default value is 20, and the value range is [1,50]. + zh_Hans: 分页大小,即单次请求所返回的数据条目数。默认值为 20,取值范围为 [1,50]。 + llm_description: 分页大小,即单次请求所返回的数据条目数。默认值为 20,取值范围为 [1,50]。 + form: form + + - name: page_token + type: string + required: false + label: + en_US: Page Token + zh_Hans: 分页标记 + human_description: + en_US: The pagination token. Leave it blank for the first request, indicating to start traversing from the beginning; when the pagination query result has more items, a new page_token will be returned simultaneously, which can be used to obtain the query result in the next traversal. + zh_Hans: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + llm_description: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_message/tools/get_thread_messages.py b/api/core/tools/provider/builtin/feishu_message/tools/get_thread_messages.py new file mode 100644 index 0000000000000000000000000000000000000000..3b14f46e0048a8e33fe68b027d6571f4d98625be --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_message/tools/get_thread_messages.py @@ -0,0 +1,21 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class GetChatMessagesTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + container_id = tool_parameters.get("container_id") + page_token = tool_parameters.get("page_token") + sort_type = tool_parameters.get("sort_type", "ByCreateTimeAsc") + page_size = tool_parameters.get("page_size", 20) + + res = client.get_thread_messages(container_id, page_token, sort_type, page_size) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_message/tools/get_thread_messages.yaml b/api/core/tools/provider/builtin/feishu_message/tools/get_thread_messages.yaml new file mode 100644 index 0000000000000000000000000000000000000000..85a138292f6203e6a3bbbec0d6320158f8694997 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_message/tools/get_thread_messages.yaml @@ -0,0 +1,72 @@ +identity: + name: get_thread_messages + author: Doug Lea + label: + en_US: Get Thread Messages + zh_Hans: 获取指定话题的消息历史 +description: + human: + en_US: Get Thread Messages + zh_Hans: 获取指定话题的消息历史 + llm: A tool for getting chat messages from specific threads.(获取指定话题的消息历史) +parameters: + - name: container_id + type: string + required: true + label: + en_US: Thread Id + zh_Hans: 话题 ID + human_description: + en_US: The ID of the thread. Refer to the thread overview on how to obtain the thread_id. https://open.feishu.cn/document/uAjLw4CM/ukTMukTMukTM/reference/im-v1/message/thread-introduction + zh_Hans: 话题 ID,获取方式参见话题概述的如何获取 thread_id 章节。https://open.feishu.cn/document/uAjLw4CM/ukTMukTMukTM/reference/im-v1/message/thread-introduction + llm_description: 话题 ID,获取方式参见话题概述的如何获取 thread_id 章节。https://open.feishu.cn/document/uAjLw4CM/ukTMukTMukTM/reference/im-v1/message/thread-introduction + form: llm + + - name: sort_type + type: select + required: false + options: + - value: ByCreateTimeAsc + label: + en_US: ByCreateTimeAsc + zh_Hans: ByCreateTimeAsc + - value: ByCreateTimeDesc + label: + en_US: ByCreateTimeDesc + zh_Hans: ByCreateTimeDesc + default: "ByCreateTimeAsc" + label: + en_US: Sort Type + zh_Hans: 排序方式 + human_description: + en_US: | + The message sorting method. Optional values are ByCreateTimeAsc: sorted in ascending order by message creation time; ByCreateTimeDesc: sorted in descending order by message creation time. The default value is ByCreateTimeAsc. Note: When using page_token for pagination requests, the sorting method (sort_type) is consistent with the first request and cannot be changed midway. + zh_Hans: | + 消息排序方式,可选值有 ByCreateTimeAsc:按消息创建时间升序排列;ByCreateTimeDesc:按消息创建时间降序排列。默认值为:ByCreateTimeAsc。注意:使用 page_token 分页请求时,排序方式(sort_type)均与第一次请求一致,不支持中途改换排序方式。 + llm_description: 消息排序方式,可选值有 ByCreateTimeAsc:按消息创建时间升序排列;ByCreateTimeDesc:按消息创建时间降序排列。默认值为:ByCreateTimeAsc。注意:使用 page_token 分页请求时,排序方式(sort_type)均与第一次请求一致,不支持中途改换排序方式。 + form: form + + - name: page_size + type: number + required: false + default: 20 + label: + en_US: Page Size + zh_Hans: 分页大小 + human_description: + en_US: The page size, i.e., the number of data entries returned in a single request. The default value is 20, and the value range is [1,50]. + zh_Hans: 分页大小,即单次请求所返回的数据条目数。默认值为 20,取值范围为 [1,50]。 + llm_description: 分页大小,即单次请求所返回的数据条目数。默认值为 20,取值范围为 [1,50]。 + form: form + + - name: page_token + type: string + required: false + label: + en_US: Page Token + zh_Hans: 分页标记 + human_description: + en_US: The pagination token. Leave it blank for the first request, indicating to start traversing from the beginning; when the pagination query result has more items, a new page_token will be returned simultaneously, which can be used to obtain the query result in the next traversal. + zh_Hans: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + llm_description: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_message/tools/send_bot_message.py b/api/core/tools/provider/builtin/feishu_message/tools/send_bot_message.py new file mode 100644 index 0000000000000000000000000000000000000000..1dd315d0e293a00535ec89c4b02b412311c0a37c --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_message/tools/send_bot_message.py @@ -0,0 +1,20 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class SendBotMessageTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + receive_id_type = tool_parameters.get("receive_id_type") + receive_id = tool_parameters.get("receive_id") + msg_type = tool_parameters.get("msg_type") + content = tool_parameters.get("content") + + res = client.send_bot_message(receive_id_type, receive_id, msg_type, content) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_message/tools/send_bot_message.yaml b/api/core/tools/provider/builtin/feishu_message/tools/send_bot_message.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4f7f65a8a74fc0ebf9abdf9fc346728c281ddfba --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_message/tools/send_bot_message.yaml @@ -0,0 +1,125 @@ +identity: + name: send_bot_message + author: Doug Lea + label: + en_US: Send Bot Message + zh_Hans: 发送飞书应用消息 +description: + human: + en_US: Send bot message + zh_Hans: 发送飞书应用消息 + llm: A tool for sending Feishu application messages. +parameters: + - name: receive_id + type: string + required: true + label: + en_US: receive_id + zh_Hans: 消息接收者的 ID + human_description: + en_US: The ID of the message receiver, the ID type is consistent with the value of the query parameter receive_id_type. + zh_Hans: 消息接收者的 ID,ID 类型与查询参数 receive_id_type 的取值一致。 + llm_description: 消息接收者的 ID,ID 类型与查询参数 receive_id_type 的取值一致。 + form: llm + + - name: receive_id_type + type: select + required: true + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + - value: email + label: + en_US: email + zh_Hans: email + - value: chat_id + label: + en_US: chat_id + zh_Hans: chat_id + label: + en_US: receive_id_type + zh_Hans: 消息接收者的 ID 类型 + human_description: + en_US: The ID type of the message receiver, optional values are open_id, union_id, user_id, email, chat_id, with a default value of open_id. + zh_Hans: 消息接收者的 ID 类型,可选值有 open_id、union_id、user_id、email、chat_id,默认值为 open_id。 + llm_description: 消息接收者的 ID 类型,可选值有 open_id、union_id、user_id、email、chat_id,默认值为 open_id。 + form: form + + - name: msg_type + type: select + required: true + options: + - value: text + label: + en_US: text + zh_Hans: 文本 + - value: interactive + label: + en_US: interactive + zh_Hans: 卡片 + - value: post + label: + en_US: post + zh_Hans: 富文本 + - value: image + label: + en_US: image + zh_Hans: 图片 + - value: file + label: + en_US: file + zh_Hans: 文件 + - value: audio + label: + en_US: audio + zh_Hans: 语音 + - value: media + label: + en_US: media + zh_Hans: 视频 + - value: sticker + label: + en_US: sticker + zh_Hans: 表情包 + - value: share_chat + label: + en_US: share_chat + zh_Hans: 分享群名片 + - value: share_user + label: + en_US: share_user + zh_Hans: 分享个人名片 + - value: system + label: + en_US: system + zh_Hans: 系统消息 + label: + en_US: msg_type + zh_Hans: 消息类型 + human_description: + en_US: Message type. Optional values are text, post, image, file, audio, media, sticker, interactive, share_chat, share_user, system. For detailed introduction of different message types, refer to the message content(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json). + zh_Hans: 消息类型。可选值有:text、post、image、file、audio、media、sticker、interactive、share_chat、share_user、system。不同消息类型的详细介绍,参见发送消息内容(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json)。 + llm_description: 消息类型。可选值有:text、post、image、file、audio、media、sticker、interactive、share_chat、share_user、system。不同消息类型的详细介绍,参见发送消息内容(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json)。 + form: form + + - name: content + type: string + required: true + label: + en_US: content + zh_Hans: 消息内容 + human_description: + en_US: Message content, a JSON structure serialized string. The value of this parameter corresponds to msg_type. For example, if msg_type is text, this parameter needs to pass in text type content. To understand the format and usage limitations of different message types, refer to the message content(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json). + zh_Hans: 消息内容,JSON 结构序列化后的字符串。该参数的取值与 msg_type 对应,例如 msg_type 取值为 text,则该参数需要传入文本类型的内容。了解不同类型的消息内容格式、使用限制,可参见发送消息内容(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json)。 + llm_description: 消息内容,JSON 结构序列化后的字符串。该参数的取值与 msg_type 对应,例如 msg_type 取值为 text,则该参数需要传入文本类型的内容。了解不同类型的消息内容格式、使用限制,可参见发送消息内容(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json)。 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_message/tools/send_webhook_message.py b/api/core/tools/provider/builtin/feishu_message/tools/send_webhook_message.py new file mode 100644 index 0000000000000000000000000000000000000000..44e70e0a15b64d03048f1abbcfceda64cbdaf541 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_message/tools/send_webhook_message.py @@ -0,0 +1,19 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class SendWebhookMessageTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + webhook = tool_parameters.get("webhook") + msg_type = tool_parameters.get("msg_type") + content = tool_parameters.get("content") + + res = client.send_webhook_message(webhook, msg_type, content) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_message/tools/send_webhook_message.yaml b/api/core/tools/provider/builtin/feishu_message/tools/send_webhook_message.yaml new file mode 100644 index 0000000000000000000000000000000000000000..eeeae8b29cd9350f2e873be58dfa243fdab016e1 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_message/tools/send_webhook_message.yaml @@ -0,0 +1,68 @@ +identity: + name: send_webhook_message + author: Doug Lea + label: + en_US: Send Webhook Message + zh_Hans: 使用自定义机器人发送飞书消息 +description: + human: + en_US: Send webhook message + zh_Hans: 使用自定义机器人发送飞书消息 + llm: A tool for sending Lark messages using a custom robot. +parameters: + - name: webhook + type: string + required: true + label: + en_US: webhook + zh_Hans: webhook + human_description: + en_US: | + The address of the webhook, the format of the webhook address corresponding to the bot is as follows: https://open.feishu.cn/open-apis/bot/v2/hook/xxxxxxxxxxxxxxxxx. For details, please refer to: Feishu Custom Bot Usage Guide(https://open.larkoffice.com/document/client-docs/bot-v3/add-custom-bot) + zh_Hans: | + webhook 的地址,机器人对应的 webhook 地址格式如下: https://open.feishu.cn/open-apis/bot/v2/hook/xxxxxxxxxxxxxxxxx,详情可参考: 飞书自定义机器人使用指南(https://open.larkoffice.com/document/client-docs/bot-v3/add-custom-bot) + llm_description: | + webhook 的地址,机器人对应的 webhook 地址格式如下: https://open.feishu.cn/open-apis/bot/v2/hook/xxxxxxxxxxxxxxxxx,详情可参考: 飞书自定义机器人使用指南(https://open.larkoffice.com/document/client-docs/bot-v3/add-custom-bot) + form: llm + + - name: msg_type + type: select + required: true + options: + - value: text + label: + en_US: text + zh_Hans: 文本 + - value: interactive + label: + en_US: interactive + zh_Hans: 卡片 + - value: image + label: + en_US: image + zh_Hans: 图片 + - value: share_chat + label: + en_US: share_chat + zh_Hans: 分享群名片 + label: + en_US: msg_type + zh_Hans: 消息类型 + human_description: + en_US: Message type. Optional values are text, image, interactive, share_chat. For detailed introduction of different message types, refer to the message content(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json). + zh_Hans: 消息类型。可选值有:text、image、interactive、share_chat。不同消息类型的详细介绍,参见发送消息内容(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json)。 + llm_description: 消息类型。可选值有:text、image、interactive、share_chat。不同消息类型的详细介绍,参见发送消息内容(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json)。 + form: form + + + - name: content + type: string + required: true + label: + en_US: content + zh_Hans: 消息内容 + human_description: + en_US: Message content, a JSON structure serialized string. The value of this parameter corresponds to msg_type. For example, if msg_type is text, this parameter needs to pass in text type content. To understand the format and usage limitations of different message types, refer to the message content(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json). + zh_Hans: 消息内容,JSON 结构序列化后的字符串。该参数的取值与 msg_type 对应,例如 msg_type 取值为 text,则该参数需要传入文本类型的内容。了解不同类型的消息内容格式、使用限制,可参见发送消息内容(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json)。 + llm_description: 消息内容,JSON 结构序列化后的字符串。该参数的取值与 msg_type 对应,例如 msg_type 取值为 text,则该参数需要传入文本类型的内容。了解不同类型的消息内容格式、使用限制,可参见发送消息内容(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json)。 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_spreadsheet/_assets/icon.png b/api/core/tools/provider/builtin/feishu_spreadsheet/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..258b361261d4e3366251613141efaf200cd492db Binary files /dev/null and b/api/core/tools/provider/builtin/feishu_spreadsheet/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/feishu_spreadsheet/feishu_spreadsheet.py b/api/core/tools/provider/builtin/feishu_spreadsheet/feishu_spreadsheet.py new file mode 100644 index 0000000000000000000000000000000000000000..a3b54737691c9cec454b239a7feedfee4327f282 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_spreadsheet/feishu_spreadsheet.py @@ -0,0 +1,7 @@ +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController +from core.tools.utils.feishu_api_utils import auth + + +class FeishuMessageProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + auth(credentials) diff --git a/api/core/tools/provider/builtin/feishu_spreadsheet/feishu_spreadsheet.yaml b/api/core/tools/provider/builtin/feishu_spreadsheet/feishu_spreadsheet.yaml new file mode 100644 index 0000000000000000000000000000000000000000..29e448d730f745e719e2ffcd04befab39c0913a5 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_spreadsheet/feishu_spreadsheet.yaml @@ -0,0 +1,36 @@ +identity: + author: Doug Lea + name: feishu_spreadsheet + label: + en_US: Feishu Spreadsheet + zh_Hans: 飞书电子表格 + description: + en_US: | + Feishu Spreadsheet, requires the following permissions: sheets:spreadsheet. + zh_Hans: | + 飞书电子表格,需要开通以下权限: sheets:spreadsheet。 + icon: icon.png + tags: + - social + - productivity +credentials_for_provider: + app_id: + type: text-input + required: true + label: + en_US: APP ID + placeholder: + en_US: Please input your feishu app id + zh_Hans: 请输入你的飞书 app id + help: + en_US: Get your app_id and app_secret from Feishu + zh_Hans: 从飞书获取您的 app_id 和 app_secret + url: https://open.larkoffice.com/app + app_secret: + type: secret-input + required: true + label: + en_US: APP Secret + placeholder: + en_US: Please input your app secret + zh_Hans: 请输入你的飞书 app secret diff --git a/api/core/tools/provider/builtin/feishu_spreadsheet/tools/add_cols.py b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/add_cols.py new file mode 100644 index 0000000000000000000000000000000000000000..44d062f9bdded2ce2ed0e9a6d5653e5d6d0c5ae5 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/add_cols.py @@ -0,0 +1,22 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class AddColsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + spreadsheet_token = tool_parameters.get("spreadsheet_token") + sheet_id = tool_parameters.get("sheet_id") + sheet_name = tool_parameters.get("sheet_name") + length = tool_parameters.get("length") + values = tool_parameters.get("values") + + res = client.add_cols(spreadsheet_token, sheet_id, sheet_name, length, values) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_spreadsheet/tools/add_cols.yaml b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/add_cols.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b73335f405c20c7ac3405552669ec34fc0cd4754 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/add_cols.yaml @@ -0,0 +1,72 @@ +identity: + name: add_cols + author: Doug Lea + label: + en_US: Add Cols + zh_Hans: 新增多列至工作表最后 +description: + human: + en_US: Add Cols + zh_Hans: 新增多列至工作表最后 + llm: A tool for adding multiple columns to the end of a spreadsheet. (新增多列至工作表最后) +parameters: + - name: spreadsheet_token + type: string + required: true + label: + en_US: spreadsheet_token + zh_Hans: 电子表格 token + human_description: + en_US: Spreadsheet token, supports input of spreadsheet URL. + zh_Hans: 电子表格 token,支持输入电子表格 url。 + llm_description: 电子表格 token,支持输入电子表格 url。 + form: llm + + - name: sheet_id + type: string + required: false + label: + en_US: sheet_id + zh_Hans: 工作表 ID + human_description: + en_US: Sheet ID, either sheet_id or sheet_name must be filled. + zh_Hans: 工作表 ID,与 sheet_name 二者其一必填。 + llm_description: 工作表 ID,与 sheet_name 二者其一必填。 + form: llm + + - name: sheet_name + type: string + required: false + label: + en_US: sheet_name + zh_Hans: 工作表名称 + human_description: + en_US: Sheet name, either sheet_id or sheet_name must be filled. + zh_Hans: 工作表名称,与 sheet_id 二者其一必填。 + llm_description: 工作表名称,与 sheet_id 二者其一必填。 + form: llm + + - name: length + type: number + required: true + label: + en_US: length + zh_Hans: 要增加的列数 + human_description: + en_US: Number of columns to add, range (0-5000]. + zh_Hans: 要增加的列数,范围(0-5000]。 + llm_description: 要增加的列数,范围(0-5000]。 + form: form + + - name: values + type: string + required: false + label: + en_US: values + zh_Hans: 新增列的单元格内容 + human_description: + en_US: | + Content of the new columns, array of objects in string format, each array represents a row of table data, format like: [ [ "ID","Name","Age" ],[ 1,"Zhang San",10 ],[ 2,"Li Si",11 ] ]. + zh_Hans: 新增列的单元格内容,数组对象字符串,每个数组一行表格数据,格式:[["编号","姓名","年龄"],[1,"张三",10],[2,"李四",11]]。 + llm_description: 新增列的单元格内容,数组对象字符串,每个数组一行表格数据,格式:[["编号","姓名","年龄"],[1,"张三",10],[2,"李四",11]]。 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_spreadsheet/tools/add_rows.py b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/add_rows.py new file mode 100644 index 0000000000000000000000000000000000000000..3a85b7b46ccb93eb3f5440f92fc9ea674dfc6bea --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/add_rows.py @@ -0,0 +1,22 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class AddRowsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + spreadsheet_token = tool_parameters.get("spreadsheet_token") + sheet_id = tool_parameters.get("sheet_id") + sheet_name = tool_parameters.get("sheet_name") + length = tool_parameters.get("length") + values = tool_parameters.get("values") + + res = client.add_rows(spreadsheet_token, sheet_id, sheet_name, length, values) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_spreadsheet/tools/add_rows.yaml b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/add_rows.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6bce305b9825ec8c826b54edbd94bbc040c75bcc --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/add_rows.yaml @@ -0,0 +1,72 @@ +identity: + name: add_rows + author: Doug Lea + label: + en_US: Add Rows + zh_Hans: 新增多行至工作表最后 +description: + human: + en_US: Add Rows + zh_Hans: 新增多行至工作表最后 + llm: A tool for adding multiple rows to the end of a spreadsheet. (新增多行至工作表最后) +parameters: + - name: spreadsheet_token + type: string + required: true + label: + en_US: spreadsheet_token + zh_Hans: 电子表格 token + human_description: + en_US: Spreadsheet token, supports input of spreadsheet URL. + zh_Hans: 电子表格 token,支持输入电子表格 url。 + llm_description: 电子表格 token,支持输入电子表格 url。 + form: llm + + - name: sheet_id + type: string + required: false + label: + en_US: sheet_id + zh_Hans: 工作表 ID + human_description: + en_US: Sheet ID, either sheet_id or sheet_name must be filled. + zh_Hans: 工作表 ID,与 sheet_name 二者其一必填。 + llm_description: 工作表 ID,与 sheet_name 二者其一必填。 + form: llm + + - name: sheet_name + type: string + required: false + label: + en_US: sheet_name + zh_Hans: 工作表名称 + human_description: + en_US: Sheet name, either sheet_id or sheet_name must be filled. + zh_Hans: 工作表名称,与 sheet_id 二者其一必填。 + llm_description: 工作表名称,与 sheet_id 二者其一必填。 + form: llm + + - name: length + type: number + required: true + label: + en_US: length + zh_Hans: 要增加行数 + human_description: + en_US: Number of rows to add, range (0-5000]. + zh_Hans: 要增加行数,范围(0-5000]。 + llm_description: 要增加行数,范围(0-5000]。 + form: form + + - name: values + type: string + required: false + label: + en_US: values + zh_Hans: 新增行的表格内容 + human_description: + en_US: | + Content of the new rows, array of objects in string format, each array represents a row of table data, format like: [ [ "ID","Name","Age" ],[ 1,"Zhang San",10 ],[ 2,"Li Si",11 ] ]. + zh_Hans: 新增行的表格内容,数组对象字符串,每个数组一行表格数据,格式,如:[["编号","姓名","年龄"],[1,"张三",10],[2,"李四",11]]。 + llm_description: 新增行的表格内容,数组对象字符串,每个数组一行表格数据,格式,如:[["编号","姓名","年龄"],[1,"张三",10],[2,"李四",11]]。 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_spreadsheet/tools/create_spreadsheet.py b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/create_spreadsheet.py new file mode 100644 index 0000000000000000000000000000000000000000..647364fab0a9660ec0603fd8ebd94eb9ed696171 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/create_spreadsheet.py @@ -0,0 +1,19 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class CreateSpreadsheetTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + title = tool_parameters.get("title") + folder_token = tool_parameters.get("folder_token") + + res = client.create_spreadsheet(title, folder_token) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_spreadsheet/tools/create_spreadsheet.yaml b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/create_spreadsheet.yaml new file mode 100644 index 0000000000000000000000000000000000000000..931310e63172d4227fd8663f66c68d201d73aa8f --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/create_spreadsheet.yaml @@ -0,0 +1,35 @@ +identity: + name: create_spreadsheet + author: Doug Lea + label: + en_US: Create Spreadsheet + zh_Hans: 创建电子表格 +description: + human: + en_US: Create Spreadsheet + zh_Hans: 创建电子表格 + llm: A tool for creating spreadsheets. (创建电子表格) +parameters: + - name: title + type: string + required: false + label: + en_US: Spreadsheet Title + zh_Hans: 电子表格标题 + human_description: + en_US: The title of the spreadsheet + zh_Hans: 电子表格的标题 + llm_description: 电子表格的标题 + form: llm + + - name: folder_token + type: string + required: false + label: + en_US: Folder Token + zh_Hans: 文件夹 token + human_description: + en_US: The token of the folder, supports folder URL input, e.g., https://bytedance.larkoffice.com/drive/folder/CxHEf4DCSlNkL2dUTCJcPRgentg + zh_Hans: 文件夹 token,支持文件夹 URL 输入,如:https://bytedance.larkoffice.com/drive/folder/CxHEf4DCSlNkL2dUTCJcPRgentg + llm_description: 文件夹 token,支持文件夹 URL 输入,如:https://bytedance.larkoffice.com/drive/folder/CxHEf4DCSlNkL2dUTCJcPRgentg + form: llm diff --git a/api/core/tools/provider/builtin/feishu_spreadsheet/tools/get_spreadsheet.py b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/get_spreadsheet.py new file mode 100644 index 0000000000000000000000000000000000000000..dda8c59daffabf91e398933e25fee399423ec407 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/get_spreadsheet.py @@ -0,0 +1,19 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class GetSpreadsheetTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + spreadsheet_token = tool_parameters.get("spreadsheet_token") + user_id_type = tool_parameters.get("user_id_type", "open_id") + + res = client.get_spreadsheet(spreadsheet_token, user_id_type) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_spreadsheet/tools/get_spreadsheet.yaml b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/get_spreadsheet.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c519938617ba8c331467aa6bf6c00c283e9b42be --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/get_spreadsheet.yaml @@ -0,0 +1,49 @@ +identity: + name: get_spreadsheet + author: Doug Lea + label: + en_US: Get Spreadsheet + zh_Hans: 获取电子表格信息 +description: + human: + en_US: Get Spreadsheet + zh_Hans: 获取电子表格信息 + llm: A tool for getting information from spreadsheets. (获取电子表格信息) +parameters: + - name: spreadsheet_token + type: string + required: true + label: + en_US: Spreadsheet Token + zh_Hans: 电子表格 token + human_description: + en_US: Spreadsheet token, supports input of spreadsheet URL. + zh_Hans: 电子表格 token,支持输入电子表格 URL。 + llm_description: 电子表格 token,支持输入电子表格 URL。 + form: llm + + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form diff --git a/api/core/tools/provider/builtin/feishu_spreadsheet/tools/list_spreadsheet_sheets.py b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/list_spreadsheet_sheets.py new file mode 100644 index 0000000000000000000000000000000000000000..98497791c0fa1ef310268d461ee87c76a4e1a255 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/list_spreadsheet_sheets.py @@ -0,0 +1,18 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class ListSpreadsheetSheetsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + spreadsheet_token = tool_parameters.get("spreadsheet_token") + + res = client.list_spreadsheet_sheets(spreadsheet_token) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_spreadsheet/tools/list_spreadsheet_sheets.yaml b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/list_spreadsheet_sheets.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c6a7ef45d46589178ccd4c8d2f145559aa2f5cd0 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/list_spreadsheet_sheets.yaml @@ -0,0 +1,23 @@ +identity: + name: list_spreadsheet_sheets + author: Doug Lea + label: + en_US: List Spreadsheet Sheets + zh_Hans: 列出电子表格所有工作表 +description: + human: + en_US: List Spreadsheet Sheets + zh_Hans: 列出电子表格所有工作表 + llm: A tool for listing all sheets in a spreadsheet. (列出电子表格所有工作表) +parameters: + - name: spreadsheet_token + type: string + required: true + label: + en_US: Spreadsheet Token + zh_Hans: 电子表格 token + human_description: + en_US: Spreadsheet token, supports input of spreadsheet URL. + zh_Hans: 电子表格 token,支持输入电子表格 URL。 + llm_description: 电子表格 token,支持输入电子表格 URL。 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_spreadsheet/tools/read_cols.py b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/read_cols.py new file mode 100644 index 0000000000000000000000000000000000000000..ebe3f619d091d1cc8f4e0b206fe39260640f3b5d --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/read_cols.py @@ -0,0 +1,23 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class ReadColsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + spreadsheet_token = tool_parameters.get("spreadsheet_token") + sheet_id = tool_parameters.get("sheet_id") + sheet_name = tool_parameters.get("sheet_name") + start_col = tool_parameters.get("start_col") + num_cols = tool_parameters.get("num_cols") + user_id_type = tool_parameters.get("user_id_type", "open_id") + + res = client.read_cols(spreadsheet_token, sheet_id, sheet_name, start_col, num_cols, user_id_type) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_spreadsheet/tools/read_cols.yaml b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/read_cols.yaml new file mode 100644 index 0000000000000000000000000000000000000000..34da74592d589864b4144aaba4b4f0777ea660d2 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/read_cols.yaml @@ -0,0 +1,97 @@ +identity: + name: read_cols + author: Doug Lea + label: + en_US: Read Cols + zh_Hans: 读取工作表列数据 +description: + human: + en_US: Read Cols + zh_Hans: 读取工作表列数据 + llm: A tool for reading column data from a spreadsheet. (读取工作表列数据) +parameters: + - name: spreadsheet_token + type: string + required: true + label: + en_US: spreadsheet_token + zh_Hans: 电子表格 token + human_description: + en_US: Spreadsheet token, supports input of spreadsheet URL. + zh_Hans: 电子表格 token,支持输入电子表格 url。 + llm_description: 电子表格 token,支持输入电子表格 url。 + form: llm + + - name: sheet_id + type: string + required: false + label: + en_US: sheet_id + zh_Hans: 工作表 ID + human_description: + en_US: Sheet ID, either sheet_id or sheet_name must be filled. + zh_Hans: 工作表 ID,与 sheet_name 二者其一必填。 + llm_description: 工作表 ID,与 sheet_name 二者其一必填。 + form: llm + + - name: sheet_name + type: string + required: false + label: + en_US: sheet_name + zh_Hans: 工作表名称 + human_description: + en_US: Sheet name, either sheet_id or sheet_name must be filled. + zh_Hans: 工作表名称,与 sheet_id 二者其一必填。 + llm_description: 工作表名称,与 sheet_id 二者其一必填。 + form: llm + + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form + + - name: start_col + type: number + required: false + label: + en_US: start_col + zh_Hans: 起始列号 + human_description: + en_US: Starting column number, starting from 1. + zh_Hans: 起始列号,从 1 开始。 + llm_description: 起始列号,从 1 开始。 + form: form + + - name: num_cols + type: number + required: true + label: + en_US: num_cols + zh_Hans: 读取列数 + human_description: + en_US: Number of columns to read. + zh_Hans: 读取列数 + llm_description: 读取列数 + form: form diff --git a/api/core/tools/provider/builtin/feishu_spreadsheet/tools/read_rows.py b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/read_rows.py new file mode 100644 index 0000000000000000000000000000000000000000..86b91b104b7029627f310a9e9d9d654762a57072 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/read_rows.py @@ -0,0 +1,23 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class ReadRowsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + spreadsheet_token = tool_parameters.get("spreadsheet_token") + sheet_id = tool_parameters.get("sheet_id") + sheet_name = tool_parameters.get("sheet_name") + start_row = tool_parameters.get("start_row") + num_rows = tool_parameters.get("num_rows") + user_id_type = tool_parameters.get("user_id_type", "open_id") + + res = client.read_rows(spreadsheet_token, sheet_id, sheet_name, start_row, num_rows, user_id_type) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_spreadsheet/tools/read_rows.yaml b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/read_rows.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5dfa8d5835412561565ab3211634167cab52c39b --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/read_rows.yaml @@ -0,0 +1,97 @@ +identity: + name: read_rows + author: Doug Lea + label: + en_US: Read Rows + zh_Hans: 读取工作表行数据 +description: + human: + en_US: Read Rows + zh_Hans: 读取工作表行数据 + llm: A tool for reading row data from a spreadsheet. (读取工作表行数据) +parameters: + - name: spreadsheet_token + type: string + required: true + label: + en_US: spreadsheet_token + zh_Hans: 电子表格 token + human_description: + en_US: Spreadsheet token, supports input of spreadsheet URL. + zh_Hans: 电子表格 token,支持输入电子表格 url。 + llm_description: 电子表格 token,支持输入电子表格 url。 + form: llm + + - name: sheet_id + type: string + required: false + label: + en_US: sheet_id + zh_Hans: 工作表 ID + human_description: + en_US: Sheet ID, either sheet_id or sheet_name must be filled. + zh_Hans: 工作表 ID,与 sheet_name 二者其一必填。 + llm_description: 工作表 ID,与 sheet_name 二者其一必填。 + form: llm + + - name: sheet_name + type: string + required: false + label: + en_US: sheet_name + zh_Hans: 工作表名称 + human_description: + en_US: Sheet name, either sheet_id or sheet_name must be filled. + zh_Hans: 工作表名称,与 sheet_id 二者其一必填。 + llm_description: 工作表名称,与 sheet_id 二者其一必填。 + form: llm + + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form + + - name: start_row + type: number + required: false + label: + en_US: start_row + zh_Hans: 起始行号 + human_description: + en_US: Starting row number, starting from 1. + zh_Hans: 起始行号,从 1 开始。 + llm_description: 起始行号,从 1 开始。 + form: form + + - name: num_rows + type: number + required: true + label: + en_US: num_rows + zh_Hans: 读取行数 + human_description: + en_US: Number of rows to read. + zh_Hans: 读取行数 + llm_description: 读取行数 + form: form diff --git a/api/core/tools/provider/builtin/feishu_spreadsheet/tools/read_table.py b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/read_table.py new file mode 100644 index 0000000000000000000000000000000000000000..ddd607d87838f46699f345d143d2eb282bfe17c2 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/read_table.py @@ -0,0 +1,23 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class ReadTableTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + spreadsheet_token = tool_parameters.get("spreadsheet_token") + sheet_id = tool_parameters.get("sheet_id") + sheet_name = tool_parameters.get("sheet_name") + num_range = tool_parameters.get("num_range") + query = tool_parameters.get("query") + user_id_type = tool_parameters.get("user_id_type", "open_id") + + res = client.read_table(spreadsheet_token, sheet_id, sheet_name, num_range, query, user_id_type) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_spreadsheet/tools/read_table.yaml b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/read_table.yaml new file mode 100644 index 0000000000000000000000000000000000000000..10534436d66e7a68c63fa191b17db8301ce4661e --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_spreadsheet/tools/read_table.yaml @@ -0,0 +1,122 @@ +identity: + name: read_table + author: Doug Lea + label: + en_US: Read Table + zh_Hans: 自定义读取电子表格行列数据 +description: + human: + en_US: Read Table + zh_Hans: 自定义读取电子表格行列数据 + llm: A tool for custom reading of row and column data from a spreadsheet. (自定义读取电子表格行列数据) +parameters: + - name: spreadsheet_token + type: string + required: true + label: + en_US: spreadsheet_token + zh_Hans: 电子表格 token + human_description: + en_US: Spreadsheet token, supports input of spreadsheet URL. + zh_Hans: 电子表格 token,支持输入电子表格 url。 + llm_description: 电子表格 token,支持输入电子表格 url。 + form: llm + + - name: sheet_id + type: string + required: false + label: + en_US: sheet_id + zh_Hans: 工作表 ID + human_description: + en_US: Sheet ID, either sheet_id or sheet_name must be filled. + zh_Hans: 工作表 ID,与 sheet_name 二者其一必填。 + llm_description: 工作表 ID,与 sheet_name 二者其一必填。 + form: llm + + - name: sheet_name + type: string + required: false + label: + en_US: sheet_name + zh_Hans: 工作表名称 + human_description: + en_US: Sheet name, either sheet_id or sheet_name must be filled. + zh_Hans: 工作表名称,与 sheet_id 二者其一必填。 + llm_description: 工作表名称,与 sheet_id 二者其一必填。 + form: llm + + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form + + - name: start_row + type: number + required: false + label: + en_US: start_row + zh_Hans: 起始行号 + human_description: + en_US: Starting row number, starting from 1. + zh_Hans: 起始行号,从 1 开始。 + llm_description: 起始行号,从 1 开始。 + form: form + + - name: num_rows + type: number + required: false + label: + en_US: num_rows + zh_Hans: 读取行数 + human_description: + en_US: Number of rows to read. + zh_Hans: 读取行数 + llm_description: 读取行数 + form: form + + - name: range + type: string + required: false + label: + en_US: range + zh_Hans: 取数范围 + human_description: + en_US: | + Data range, format like: A1:B2, can be empty when query=all. + zh_Hans: 取数范围,格式如:A1:B2,query=all 时可为空。 + llm_description: 取数范围,格式如:A1:B2,query=all 时可为空。 + form: llm + + - name: query + type: string + required: false + label: + en_US: query + zh_Hans: 查询 + human_description: + en_US: Pass "all" to query all data in the table, but no more than 100 columns. + zh_Hans: 传 all,表示查询表格所有数据,但最多查询 100 列数据。 + llm_description: 传 all,表示查询表格所有数据,但最多查询 100 列数据。 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_task/_assets/icon.png b/api/core/tools/provider/builtin/feishu_task/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..3485be0d0fbd85444995dca78b51634899807537 Binary files /dev/null and b/api/core/tools/provider/builtin/feishu_task/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/feishu_task/feishu_task.py b/api/core/tools/provider/builtin/feishu_task/feishu_task.py new file mode 100644 index 0000000000000000000000000000000000000000..6df05968d8f176c7d9b402a59eb34426c0a8de85 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_task/feishu_task.py @@ -0,0 +1,7 @@ +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController +from core.tools.utils.feishu_api_utils import auth + + +class FeishuTaskProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + auth(credentials) diff --git a/api/core/tools/provider/builtin/feishu_task/feishu_task.yaml b/api/core/tools/provider/builtin/feishu_task/feishu_task.yaml new file mode 100644 index 0000000000000000000000000000000000000000..88736f79a02e879e9198bc39533adf8b47e08c8c --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_task/feishu_task.yaml @@ -0,0 +1,36 @@ +identity: + author: Doug Lea + name: feishu_task + label: + en_US: Feishu Task + zh_Hans: 飞书任务 + description: + en_US: | + Feishu Task, requires the following permissions: task:task:write、contact:user.id:readonly. + zh_Hans: | + 飞书任务,需要开通以下权限: task:task:write、contact:user.id:readonly。 + icon: icon.png + tags: + - social + - productivity +credentials_for_provider: + app_id: + type: text-input + required: true + label: + en_US: APP ID + placeholder: + en_US: Please input your feishu app id + zh_Hans: 请输入你的飞书 app id + help: + en_US: Get your app_id and app_secret from Feishu + zh_Hans: 从飞书获取您的 app_id 和 app_secret + url: https://open.larkoffice.com/app + app_secret: + type: secret-input + required: true + label: + en_US: APP Secret + placeholder: + en_US: Please input your app secret + zh_Hans: 请输入你的飞书 app secret diff --git a/api/core/tools/provider/builtin/feishu_task/tools/add_members.py b/api/core/tools/provider/builtin/feishu_task/tools/add_members.py new file mode 100644 index 0000000000000000000000000000000000000000..e58ed22e0f4797c6868a0d60321b903af7a2d5fe --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_task/tools/add_members.py @@ -0,0 +1,20 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class AddMembersTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + task_guid = tool_parameters.get("task_guid") + member_phone_or_email = tool_parameters.get("member_phone_or_email") + member_role = tool_parameters.get("member_role", "follower") + + res = client.add_members(task_guid, member_phone_or_email, member_role) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_task/tools/add_members.yaml b/api/core/tools/provider/builtin/feishu_task/tools/add_members.yaml new file mode 100644 index 0000000000000000000000000000000000000000..063c0f7f04956cb10077ad372c7c5399d3f1251b --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_task/tools/add_members.yaml @@ -0,0 +1,58 @@ +identity: + name: add_members + author: Doug Lea + label: + en_US: Add Members + zh_Hans: 添加任务成员 +description: + human: + en_US: Add Members + zh_Hans: 添加任务成员 + llm: A tool for adding members to a Feishu task.(添加任务成员) +parameters: + - name: task_guid + type: string + required: true + label: + en_US: Task GUID + zh_Hans: 任务 GUID + human_description: + en_US: | + The GUID of the task to be added, supports passing either the Task ID or the Task link URL. Example of Task ID: 8b5425ec-9f2a-43bd-a3ab-01912f50282b; Example of Task link URL: https://applink.feishu-pre.net/client/todo/detail?guid=8c6bf822-e4da-449a-b82a-dc44020f9be9&suite_entity_num=t21587362 + zh_Hans: 要添加的任务的 GUID,支持传任务 ID 和任务链接 URL。任务 ID 示例:8b5425ec-9f2a-43bd-a3ab-01912f50282b;任务链接 URL 示例:https://applink.feishu-pre.net/client/todo/detail?guid=8c6bf822-e4da-449a-b82a-dc44020f9be9&suite_entity_num=t21587362 + llm_description: 要添加的任务的 GUID,支持传任务 ID 和任务链接 URL。任务 ID 示例:8b5425ec-9f2a-43bd-a3ab-01912f50282b;任务链接 URL 示例:https://applink.feishu-pre.net/client/todo/detail?guid=8c6bf822-e4da-449a-b82a-dc44020f9be9&suite_entity_num=t21587362 + form: llm + + - name: member_phone_or_email + type: string + required: true + label: + en_US: Task Member Phone Or Email + zh_Hans: 任务成员的电话或邮箱 + human_description: + en_US: A list of member emails or phone numbers, separated by commas. + zh_Hans: 任务成员邮箱或者手机号列表,使用逗号分隔。 + llm_description: 任务成员邮箱或者手机号列表,使用逗号分隔。 + form: llm + + - name: member_role + type: select + required: true + options: + - value: assignee + label: + en_US: assignee + zh_Hans: 负责人 + - value: follower + label: + en_US: follower + zh_Hans: 关注人 + default: "follower" + label: + en_US: member_role + zh_Hans: 成员的角色 + human_description: + en_US: Member role, optional values are "assignee" (responsible person) and "follower" (observer), with a default value of "assignee". + zh_Hans: 成员的角色,可选值有 "assignee"(负责人)和 "follower"(关注人),默认值为 "assignee"。 + llm_description: 成员的角色,可选值有 "assignee"(负责人)和 "follower"(关注人),默认值为 "assignee"。 + form: form diff --git a/api/core/tools/provider/builtin/feishu_task/tools/create_task.py b/api/core/tools/provider/builtin/feishu_task/tools/create_task.py new file mode 100644 index 0000000000000000000000000000000000000000..96cdcd71f6d2ec37bd68dd792cced40d4c02b08c --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_task/tools/create_task.py @@ -0,0 +1,22 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class CreateTaskTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + summary = tool_parameters.get("summary") + start_time = tool_parameters.get("start_time") + end_time = tool_parameters.get("end_time") + completed_time = tool_parameters.get("completed_time") + description = tool_parameters.get("description") + + res = client.create_task(summary, start_time, end_time, completed_time, description) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_task/tools/create_task.yaml b/api/core/tools/provider/builtin/feishu_task/tools/create_task.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7eb4af168bf740e99d684bff7ff67494582a31e8 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_task/tools/create_task.yaml @@ -0,0 +1,74 @@ +identity: + name: create_task + author: Doug Lea + label: + en_US: Create Task + zh_Hans: 创建飞书任务 +description: + human: + en_US: Create Feishu Task + zh_Hans: 创建飞书任务 + llm: A tool for creating tasks in Feishu.(创建飞书任务) +parameters: + - name: summary + type: string + required: true + label: + en_US: Task Title + zh_Hans: 任务标题 + human_description: + en_US: The title of the task. + zh_Hans: 任务标题 + llm_description: 任务标题 + form: llm + + - name: description + type: string + required: false + label: + en_US: Task Description + zh_Hans: 任务备注 + human_description: + en_US: The description or notes for the task. + zh_Hans: 任务备注 + llm_description: 任务备注 + form: llm + + - name: start_time + type: string + required: false + label: + en_US: Start Time + zh_Hans: 任务开始时间 + human_description: + en_US: | + The start time of the task, in the format: 2006-01-02 15:04:05 + zh_Hans: 任务开始时间,格式为:2006-01-02 15:04:05 + llm_description: 任务开始时间,格式为:2006-01-02 15:04:05 + form: llm + + - name: end_time + type: string + required: false + label: + en_US: End Time + zh_Hans: 任务结束时间 + human_description: + en_US: | + The end time of the task, in the format: 2006-01-02 15:04:05 + zh_Hans: 任务结束时间,格式为:2006-01-02 15:04:05 + llm_description: 任务结束时间,格式为:2006-01-02 15:04:05 + form: llm + + - name: completed_time + type: string + required: false + label: + en_US: Completed Time + zh_Hans: 任务完成时间 + human_description: + en_US: | + The completion time of the task, in the format: 2006-01-02 15:04:05. Leave empty to create an incomplete task; fill in a specific time to create a completed task. + zh_Hans: 任务完成时间,格式为:2006-01-02 15:04:05,不填写表示创建一个未完成任务;填写一个具体的时间表示创建一个已完成任务。 + llm_description: 任务完成时间,格式为:2006-01-02 15:04:05,不填写表示创建一个未完成任务;填写一个具体的时间表示创建一个已完成任务。 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_task/tools/delete_task.py b/api/core/tools/provider/builtin/feishu_task/tools/delete_task.py new file mode 100644 index 0000000000000000000000000000000000000000..dee036fee5203afd424d445eec968057f4be574f --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_task/tools/delete_task.py @@ -0,0 +1,18 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class UpdateTaskTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + task_guid = tool_parameters.get("task_guid") + + res = client.delete_task(task_guid) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_task/tools/delete_task.yaml b/api/core/tools/provider/builtin/feishu_task/tools/delete_task.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d3f9741367662431c5d5f7e5b4bb984e2ea9d8b8 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_task/tools/delete_task.yaml @@ -0,0 +1,24 @@ +identity: + name: delete_task + author: Doug Lea + label: + en_US: Delete Task + zh_Hans: 删除飞书任务 +description: + human: + en_US: Delete Task + zh_Hans: 删除飞书任务 + llm: A tool for deleting tasks in Feishu.(删除飞书任务) +parameters: + - name: task_guid + type: string + required: true + label: + en_US: Task GUID + zh_Hans: 任务 GUID + human_description: + en_US: | + The GUID of the task to be deleted, supports passing either the Task ID or the Task link URL. Example of Task ID: 8b5425ec-9f2a-43bd-a3ab-01912f50282b; Example of Task link URL: https://applink.feishu-pre.net/client/todo/detail?guid=8c6bf822-e4da-449a-b82a-dc44020f9be9&suite_entity_num=t21587362 + zh_Hans: 要删除的任务的 GUID,支持传任务 ID 和任务链接 URL。任务 ID 示例:8b5425ec-9f2a-43bd-a3ab-01912f50282b;任务链接 URL 示例:https://applink.feishu-pre.net/client/todo/detail?guid=8c6bf822-e4da-449a-b82a-dc44020f9be9&suite_entity_num=t21587362 + llm_description: 要删除的任务的 GUID,支持传任务 ID 和任务链接 URL。任务 ID 示例:8b5425ec-9f2a-43bd-a3ab-01912f50282b;任务链接 URL 示例:https://applink.feishu-pre.net/client/todo/detail?guid=8c6bf822-e4da-449a-b82a-dc44020f9be9&suite_entity_num=t21587362 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_task/tools/update_task.py b/api/core/tools/provider/builtin/feishu_task/tools/update_task.py new file mode 100644 index 0000000000000000000000000000000000000000..4a48cd283abf1df2fdf280f9ff7395921fe825ec --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_task/tools/update_task.py @@ -0,0 +1,23 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class UpdateTaskTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + task_guid = tool_parameters.get("task_guid") + summary = tool_parameters.get("summary") + start_time = tool_parameters.get("start_time") + end_time = tool_parameters.get("end_time") + completed_time = tool_parameters.get("completed_time") + description = tool_parameters.get("description") + + res = client.update_task(task_guid, summary, start_time, end_time, completed_time, description) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_task/tools/update_task.yaml b/api/core/tools/provider/builtin/feishu_task/tools/update_task.yaml new file mode 100644 index 0000000000000000000000000000000000000000..83c9bcb1c443ac595f1f0893fd10eae6959a7602 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_task/tools/update_task.yaml @@ -0,0 +1,89 @@ +identity: + name: update_task + author: Doug Lea + label: + en_US: Update Task + zh_Hans: 更新飞书任务 +description: + human: + en_US: Update Feishu Task + zh_Hans: 更新飞书任务 + llm: A tool for updating tasks in Feishu.(更新飞书任务) +parameters: + - name: task_guid + type: string + required: true + label: + en_US: Task GUID + zh_Hans: 任务 GUID + human_description: + en_US: | + The task ID, supports inputting either the Task ID or the Task link URL. Example of Task ID: 42cad8a0-f8c8-4344-9be2-d1d7e8e91b64; Example of Task link URL: https://applink.feishu-pre.net/client/todo/detail?guid=42cad8a0-f8c8-4344-9be2-d1d7e8e91b64&suite_entity_num=t21700217 + zh_Hans: | + 任务ID,支持传入任务 ID 和任务链接 URL。任务 ID 示例: 42cad8a0-f8c8-4344-9be2-d1d7e8e91b64;任务链接 URL 示例: https://applink.feishu-pre.net/client/todo/detail?guid=42cad8a0-f8c8-4344-9be2-d1d7e8e91b64&suite_entity_num=t21700217 + llm_description: | + 任务ID,支持传入任务 ID 和任务链接 URL。任务 ID 示例: 42cad8a0-f8c8-4344-9be2-d1d7e8e91b64;任务链接 URL 示例: https://applink.feishu-pre.net/client/todo/detail?guid=42cad8a0-f8c8-4344-9be2-d1d7e8e91b64&suite_entity_num=t21700217 + form: llm + + - name: summary + type: string + required: true + label: + en_US: Task Title + zh_Hans: 任务标题 + human_description: + en_US: The title of the task. + zh_Hans: 任务标题 + llm_description: 任务标题 + form: llm + + - name: description + type: string + required: false + label: + en_US: Task Description + zh_Hans: 任务备注 + human_description: + en_US: The description or notes for the task. + zh_Hans: 任务备注 + llm_description: 任务备注 + form: llm + + - name: start_time + type: string + required: false + label: + en_US: Start Time + zh_Hans: 任务开始时间 + human_description: + en_US: | + The start time of the task, in the format: 2006-01-02 15:04:05 + zh_Hans: 任务开始时间,格式为:2006-01-02 15:04:05 + llm_description: 任务开始时间,格式为:2006-01-02 15:04:05 + form: llm + + - name: end_time + type: string + required: false + label: + en_US: End Time + zh_Hans: 任务结束时间 + human_description: + en_US: | + The end time of the task, in the format: 2006-01-02 15:04:05 + zh_Hans: 任务结束时间,格式为:2006-01-02 15:04:05 + llm_description: 任务结束时间,格式为:2006-01-02 15:04:05 + form: llm + + - name: completed_time + type: string + required: false + label: + en_US: Completed Time + zh_Hans: 任务完成时间 + human_description: + en_US: | + The completion time of the task, in the format: 2006-01-02 15:04:05 + zh_Hans: 任务完成时间,格式为:2006-01-02 15:04:05 + llm_description: 任务完成时间,格式为:2006-01-02 15:04:05 + form: llm diff --git a/api/core/tools/provider/builtin/feishu_wiki/_assets/icon.png b/api/core/tools/provider/builtin/feishu_wiki/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..878672c9ae5a511f246d3211863ec786596481eb Binary files /dev/null and b/api/core/tools/provider/builtin/feishu_wiki/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/feishu_wiki/feishu_wiki.py b/api/core/tools/provider/builtin/feishu_wiki/feishu_wiki.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5fccb1a31d0df21f5139455bb2e4a728d63343 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_wiki/feishu_wiki.py @@ -0,0 +1,7 @@ +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController +from core.tools.utils.feishu_api_utils import auth + + +class FeishuWikiProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + auth(credentials) diff --git a/api/core/tools/provider/builtin/feishu_wiki/feishu_wiki.yaml b/api/core/tools/provider/builtin/feishu_wiki/feishu_wiki.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1fb5f71cbc51692c5c9a67413162139c9a713f42 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_wiki/feishu_wiki.yaml @@ -0,0 +1,36 @@ +identity: + author: Doug Lea + name: feishu_wiki + label: + en_US: Feishu Wiki + zh_Hans: 飞书知识库 + description: + en_US: | + Feishu Wiki, requires the following permissions: wiki:wiki:readonly. + zh_Hans: | + 飞书知识库,需要开通以下权限: wiki:wiki:readonly。 + icon: icon.png + tags: + - social + - productivity +credentials_for_provider: + app_id: + type: text-input + required: true + label: + en_US: APP ID + placeholder: + en_US: Please input your feishu app id + zh_Hans: 请输入你的飞书 app id + help: + en_US: Get your app_id and app_secret from Feishu + zh_Hans: 从飞书获取您的 app_id 和 app_secret + url: https://open.larkoffice.com/app + app_secret: + type: secret-input + required: true + label: + en_US: APP Secret + placeholder: + en_US: Please input your app secret + zh_Hans: 请输入你的飞书 app secret diff --git a/api/core/tools/provider/builtin/feishu_wiki/tools/get_wiki_nodes.py b/api/core/tools/provider/builtin/feishu_wiki/tools/get_wiki_nodes.py new file mode 100644 index 0000000000000000000000000000000000000000..374b4c9a7d14923251a1d0fd62298d3465d34818 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_wiki/tools/get_wiki_nodes.py @@ -0,0 +1,21 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.feishu_api_utils import FeishuRequest + + +class GetWikiNodesTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = FeishuRequest(app_id, app_secret) + + space_id = tool_parameters.get("space_id") + parent_node_token = tool_parameters.get("parent_node_token") + page_token = tool_parameters.get("page_token") + page_size = tool_parameters.get("page_size") + + res = client.get_wiki_nodes(space_id, parent_node_token, page_token, page_size) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/feishu_wiki/tools/get_wiki_nodes.yaml b/api/core/tools/provider/builtin/feishu_wiki/tools/get_wiki_nodes.yaml new file mode 100644 index 0000000000000000000000000000000000000000..74d51e7bcbc32a57d93286b81699e8b5fd2fa402 --- /dev/null +++ b/api/core/tools/provider/builtin/feishu_wiki/tools/get_wiki_nodes.yaml @@ -0,0 +1,63 @@ +identity: + name: get_wiki_nodes + author: Doug Lea + label: + en_US: Get Wiki Nodes + zh_Hans: 获取知识空间子节点列表 +description: + human: + en_US: | + Get the list of child nodes in Wiki, make sure the app/bot is a member of the wiki space. See How to add an app as a wiki base administrator (member). https://open.feishu.cn/document/server-docs/docs/wiki-v2/wiki-qa + zh_Hans: | + 获取知识库全部子节点列表,请确保应用/机器人为知识空间成员。参阅如何将应用添加为知识库管理员(成员)。https://open.feishu.cn/document/server-docs/docs/wiki-v2/wiki-qa + llm: A tool for getting all sub-nodes of a knowledge base.(获取知识空间子节点列表) +parameters: + - name: space_id + type: string + required: true + label: + en_US: Space Id + zh_Hans: 知识空间 ID + human_description: + en_US: | + The ID of the knowledge space. Supports space link URL, for example: https://svi136aogf123.feishu.cn/wiki/settings/7166950623940706332 + zh_Hans: 知识空间 ID,支持空间链接 URL,例如:https://svi136aogf123.feishu.cn/wiki/settings/7166950623940706332 + llm_description: 知识空间 ID,支持空间链接 URL,例如:https://svi136aogf123.feishu.cn/wiki/settings/7166950623940706332 + form: llm + + - name: page_size + type: number + required: false + default: 10 + label: + en_US: Page Size + zh_Hans: 分页大小 + human_description: + en_US: The size of each page, with a maximum value of 50. + zh_Hans: 分页大小,最大值 50。 + llm_description: 分页大小,最大值 50。 + form: form + + - name: page_token + type: string + required: false + label: + en_US: Page Token + zh_Hans: 分页标记 + human_description: + en_US: The pagination token. Leave empty for the first request to start from the beginning; if the paginated query result has more items, a new page_token will be returned, which can be used to get the next set of results. + zh_Hans: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + llm_description: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + form: llm + + - name: parent_node_token + type: string + required: false + label: + en_US: Parent Node Token + zh_Hans: 父节点 token + human_description: + en_US: The token of the parent node. + zh_Hans: 父节点 token + llm_description: 父节点 token + form: llm diff --git a/api/core/tools/provider/builtin/firecrawl/_assets/icon.svg b/api/core/tools/provider/builtin/firecrawl/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..e1e5f54117b1be9ebdc2fb293940c27286df4b4f --- /dev/null +++ b/api/core/tools/provider/builtin/firecrawl/_assets/icon.svg @@ -0,0 +1,3 @@ + + 🔥 + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/firecrawl/firecrawl.py b/api/core/tools/provider/builtin/firecrawl/firecrawl.py new file mode 100644 index 0000000000000000000000000000000000000000..01455d7206f185bbf2f370243edf8e10c10e5c5d --- /dev/null +++ b/api/core/tools/provider/builtin/firecrawl/firecrawl.py @@ -0,0 +1,14 @@ +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.firecrawl.tools.scrape import ScrapeTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class FirecrawlProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + try: + # Example validation using the ScrapeTool, only scraping title for minimize content + ScrapeTool().fork_tool_runtime(runtime={"credentials": credentials}).invoke( + user_id="", tool_parameters={"url": "https://google.com", "onlyIncludeTags": "title"} + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/firecrawl/firecrawl.yaml b/api/core/tools/provider/builtin/firecrawl/firecrawl.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a48b9d9f541eb3b0fdf2b4098c4aaf94b4e41fdd --- /dev/null +++ b/api/core/tools/provider/builtin/firecrawl/firecrawl.yaml @@ -0,0 +1,35 @@ +identity: + author: Richards Tu + name: firecrawl + label: + en_US: Firecrawl + zh_CN: Firecrawl + description: + en_US: Firecrawl API integration for web crawling and scraping. + zh_Hans: Firecrawl API 集成,用于网页爬取和数据抓取。 + icon: icon.svg + tags: + - search + - utilities +credentials_for_provider: + firecrawl_api_key: + type: secret-input + required: true + label: + en_US: Firecrawl API Key + zh_Hans: Firecrawl API 密钥 + placeholder: + en_US: Please input your Firecrawl API key + zh_Hans: 请输入您的 Firecrawl API 密钥,如果是自托管版本,可以随意填写密钥 + help: + en_US: Get your Firecrawl API key from your Firecrawl account settings.If you are using a self-hosted version, you may enter any key at your convenience. + zh_Hans: 从您的 Firecrawl 账户设置中获取 Firecrawl API 密钥。如果是自托管版本,可以随意填写密钥。 + url: https://www.firecrawl.dev/account + base_url: + type: text-input + required: false + label: + en_US: Firecrawl server's Base URL + zh_Hans: Firecrawl服务器的API URL + placeholder: + en_US: https://api.firecrawl.dev diff --git a/api/core/tools/provider/builtin/firecrawl/firecrawl_appx.py b/api/core/tools/provider/builtin/firecrawl/firecrawl_appx.py new file mode 100644 index 0000000000000000000000000000000000000000..14596bf93f493fd0553e7acbf1bc930073faea72 --- /dev/null +++ b/api/core/tools/provider/builtin/firecrawl/firecrawl_appx.py @@ -0,0 +1,122 @@ +import json +import logging +import time +from collections.abc import Mapping +from typing import Any + +import requests +from requests.exceptions import HTTPError + +logger = logging.getLogger(__name__) + + +class FirecrawlApp: + def __init__(self, api_key: str | None = None, base_url: str | None = None): + self.api_key = api_key + self.base_url = base_url or "https://api.firecrawl.dev" + if not self.api_key: + raise ValueError("API key is required") + + def _prepare_headers(self, idempotency_key: str | None = None): + headers = {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"} + if idempotency_key: + headers["Idempotency-Key"] = idempotency_key + return headers + + def _request( + self, + method: str, + url: str, + data: Mapping[str, Any] | None = None, + headers: Mapping[str, str] | None = None, + retries: int = 3, + backoff_factor: float = 0.3, + ) -> Mapping[str, Any] | None: + if not headers: + headers = self._prepare_headers() + for i in range(retries): + try: + response = requests.request(method, url, json=data, headers=headers) + return response.json() + except requests.exceptions.RequestException: + if i < retries - 1: + time.sleep(backoff_factor * (2**i)) + else: + raise + return None + + def scrape_url(self, url: str, **kwargs): + endpoint = f"{self.base_url}/v1/scrape" + data = {"url": url, **kwargs} + logger.debug(f"Sent request to {endpoint=} body={data}") + response = self._request("POST", endpoint, data) + if response is None: + raise HTTPError("Failed to scrape URL after multiple retries") + return response + + def map(self, url: str, **kwargs): + endpoint = f"{self.base_url}/v1/map" + data = {"url": url, **kwargs} + logger.debug(f"Sent request to {endpoint=} body={data}") + response = self._request("POST", endpoint, data) + if response is None: + raise HTTPError("Failed to perform map after multiple retries") + return response + + def crawl_url( + self, url: str, wait: bool = True, poll_interval: int = 5, idempotency_key: str | None = None, **kwargs + ): + endpoint = f"{self.base_url}/v1/crawl" + headers = self._prepare_headers(idempotency_key) + data = {"url": url, **kwargs} + logger.debug(f"Sent request to {endpoint=} body={data}") + response = self._request("POST", endpoint, data, headers) + if response is None: + raise HTTPError("Failed to initiate crawl after multiple retries") + elif response.get("success") == False: + raise HTTPError(f"Failed to crawl: {response.get('error')}") + job_id: str = response["id"] + if wait: + return self._monitor_job_status(job_id=job_id, poll_interval=poll_interval) + return response + + def check_crawl_status(self, job_id: str): + endpoint = f"{self.base_url}/v1/crawl/{job_id}" + response = self._request("GET", endpoint) + if response is None: + raise HTTPError(f"Failed to check status for job {job_id} after multiple retries") + return response + + def cancel_crawl_job(self, job_id: str): + endpoint = f"{self.base_url}/v1/crawl/{job_id}" + response = self._request("DELETE", endpoint) + if response is None: + raise HTTPError(f"Failed to cancel job {job_id} after multiple retries") + return response + + def _monitor_job_status(self, job_id: str, poll_interval: int): + while True: + status = self.check_crawl_status(job_id) + if status["status"] == "completed": + return status + elif status["status"] == "failed": + raise HTTPError(f"Job {job_id} failed: {status['error']}") + time.sleep(poll_interval) + + +def get_array_params(tool_parameters: dict[str, Any], key): + param = tool_parameters.get(key) + if param: + return param.split(",") + + +def get_json_params(tool_parameters: dict[str, Any], key): + param = tool_parameters.get(key) + if param: + try: + # support both single quotes and double quotes + param = param.replace("'", '"') + param = json.loads(param) + except Exception: + raise ValueError(f"Invalid {key} format.") + return param diff --git a/api/core/tools/provider/builtin/firecrawl/tools/crawl.py b/api/core/tools/provider/builtin/firecrawl/tools/crawl.py new file mode 100644 index 0000000000000000000000000000000000000000..15ab510c6c889ca2e43393bbf03a2c1f609400cb --- /dev/null +++ b/api/core/tools/provider/builtin/firecrawl/tools/crawl.py @@ -0,0 +1,45 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.firecrawl.firecrawl_appx import FirecrawlApp, get_array_params, get_json_params +from core.tools.tool.builtin_tool import BuiltinTool + + +class CrawlTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + """ + the api doc: + https://docs.firecrawl.dev/api-reference/endpoint/crawl + """ + app = FirecrawlApp( + api_key=self.runtime.credentials["firecrawl_api_key"], base_url=self.runtime.credentials["base_url"] + ) + + scrapeOptions = {} + payload = {} + + wait_for_results = tool_parameters.get("wait_for_results", True) + + payload["excludePaths"] = get_array_params(tool_parameters, "excludePaths") + payload["includePaths"] = get_array_params(tool_parameters, "includePaths") + payload["maxDepth"] = tool_parameters.get("maxDepth") + payload["ignoreSitemap"] = tool_parameters.get("ignoreSitemap", False) + payload["limit"] = tool_parameters.get("limit", 5) + payload["allowBackwardLinks"] = tool_parameters.get("allowBackwardLinks", False) + payload["allowExternalLinks"] = tool_parameters.get("allowExternalLinks", False) + payload["webhook"] = tool_parameters.get("webhook") + + scrapeOptions["formats"] = get_array_params(tool_parameters, "formats") + scrapeOptions["headers"] = get_json_params(tool_parameters, "headers") + scrapeOptions["includeTags"] = get_array_params(tool_parameters, "includeTags") + scrapeOptions["excludeTags"] = get_array_params(tool_parameters, "excludeTags") + scrapeOptions["onlyMainContent"] = tool_parameters.get("onlyMainContent", False) + scrapeOptions["waitFor"] = tool_parameters.get("waitFor", 0) + scrapeOptions = {k: v for k, v in scrapeOptions.items() if v not in (None, "")} + payload["scrapeOptions"] = scrapeOptions or None + + payload = {k: v for k, v in payload.items() if v not in (None, "")} + + crawl_result = app.crawl_url(url=tool_parameters["url"], wait=wait_for_results, **payload) + + return self.create_json_message(crawl_result) diff --git a/api/core/tools/provider/builtin/firecrawl/tools/crawl.yaml b/api/core/tools/provider/builtin/firecrawl/tools/crawl.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0d7dbcac20ea16c48114730c36df4af51222427f --- /dev/null +++ b/api/core/tools/provider/builtin/firecrawl/tools/crawl.yaml @@ -0,0 +1,200 @@ +identity: + name: crawl + author: Richards Tu + label: + en_US: Crawl + zh_Hans: 深度爬取 +description: + human: + en_US: Recursively search through a urls subdomains, and gather the content. + zh_Hans: 递归爬取一个网址的子域名,并收集内容。 + llm: This tool initiates a web crawl to extract data from a specified URL. It allows configuring crawler options such as including or excluding URL patterns, generating alt text for images using LLMs (paid plan required), limiting the maximum number of pages to crawl, and returning only the main content of the page. The tool can return either a list of crawled documents or a list of URLs based on the provided options. +parameters: + - name: url + type: string + required: true + label: + en_US: Start URL + zh_Hans: 起始URL + human_description: + en_US: The base URL to start crawling from. + zh_Hans: 要爬取网站的起始URL。 + llm_description: The URL of the website that needs to be crawled. This is a required parameter. + form: llm + - name: wait_for_results + type: boolean + default: true + label: + en_US: Wait For Results + zh_Hans: 等待爬取结果 + human_description: + en_US: If you choose not to wait, it will directly return a job ID. You can use this job ID to check the crawling results or cancel the crawling task, which is usually very useful for a large-scale crawling task. + zh_Hans: 如果选择不等待,则会直接返回一个job_id,可以通过job_id查询爬取结果或取消爬取任务,这通常对于一个大型爬取任务来说非常有用。 + form: form +############## Payload ####################### + - name: excludePaths + type: string + label: + en_US: URL patterns to exclude + zh_Hans: 要排除的URL模式 + placeholder: + en_US: Use commas to separate multiple tags + zh_Hans: 多个标签时使用半角逗号分隔 + human_description: + en_US: | + Pages matching these patterns will be skipped. Example: blog/*, about/* + zh_Hans: 匹配这些模式的页面将被跳过。示例:blog/*, about/* + form: form + - name: includePaths + type: string + required: false + label: + en_US: URL patterns to include + zh_Hans: 要包含的URL模式 + placeholder: + en_US: Use commas to separate multiple tags + zh_Hans: 多个标签时使用半角逗号分隔 + human_description: + en_US: | + Only pages matching these patterns will be crawled. Example: blog/*, about/* + zh_Hans: 只有与这些模式匹配的页面才会被爬取。示例:blog/*, about/* + form: form + - name: maxDepth + type: number + label: + en_US: Maximum crawl depth + zh_Hans: 爬取深度 + human_description: + en_US: Maximum depth to crawl relative to the entered URL. A maxDepth of 0 scrapes only the entered URL. A maxDepth of 1 scrapes the entered URL and all pages one level deep. A maxDepth of 2 scrapes the entered URL and all pages up to two levels deep. Higher values follow the same pattern. + zh_Hans: 相对于输入的URL,爬取的最大深度。maxDepth为0时,仅抓取输入的URL。maxDepth为1时,抓取输入的URL以及所有一级深层页面。maxDepth为2时,抓取输入的URL以及所有两级深层页面。更高值遵循相同模式。 + form: form + min: 0 + default: 2 + - name: ignoreSitemap + type: boolean + default: true + label: + en_US: ignore Sitemap + zh_Hans: 忽略站点地图 + human_description: + en_US: Ignore the website sitemap when crawling. + zh_Hans: 爬取时忽略网站站点地图。 + form: form + - name: limit + type: number + required: false + label: + en_US: Maximum pages to crawl + zh_Hans: 最大爬取页面数 + human_description: + en_US: Specify the maximum number of pages to crawl. The crawler will stop after reaching this limit. + zh_Hans: 指定要爬取的最大页面数。爬虫将在达到此限制后停止。 + form: form + min: 1 + default: 5 + - name: allowBackwardLinks + type: boolean + default: false + label: + en_US: allow Backward Crawling + zh_Hans: 允许向后爬取 + human_description: + en_US: Enables the crawler to navigate from a specific URL to previously linked pages. For instance, from 'example.com/product/123' back to 'example.com/product' + zh_Hans: 使爬虫能够从特定URL导航到之前链接的页面。例如,从'example.com/product/123'返回到'example.com/product' + form: form + - name: allowExternalLinks + type: boolean + default: false + label: + en_US: allow External Content Links + zh_Hans: 允许爬取外链 + human_description: + en_US: Allows the crawler to follow links to external websites. + zh_Hans: + form: form + - name: webhook + type: string + label: + en_US: webhook + human_description: + en_US: | + The URL to send the webhook to. This will trigger for crawl started (crawl.started) ,every page crawled (crawl.page) and when the crawl is completed (crawl.completed or crawl.failed). The response will be the same as the /scrape endpoint. + zh_Hans: 发送Webhook的URL。这将在开始爬取(crawl.started)、每爬取一个页面(crawl.page)以及爬取完成(crawl.completed或crawl.failed)时触发。响应将与/scrape端点相同。 + form: form +############## Scrape Options ####################### + - name: formats + type: string + label: + en_US: Formats + zh_Hans: 结果的格式 + placeholder: + en_US: Use commas to separate multiple tags + zh_Hans: 多个标签时使用半角逗号分隔 + human_description: + en_US: | + Formats to include in the output. Available options: markdown, html, rawHtml, links, screenshot + zh_Hans: | + 输出中应包含的格式。可以填入: markdown, html, rawHtml, links, screenshot + form: form + - name: headers + type: string + label: + en_US: headers + zh_Hans: 请求头 + human_description: + en_US: | + Headers to send with the request. Can be used to send cookies, user-agent, etc. Example: {"cookies": "testcookies"} + zh_Hans: | + 随请求发送的头部。可以用来发送cookies、用户代理等。示例:{"cookies": "testcookies"} + placeholder: + en_US: Please enter an object that can be serialized in JSON + zh_Hans: 请输入可以json序列化的对象 + form: form + - name: includeTags + type: string + label: + en_US: Include Tags + zh_Hans: 仅抓取这些标签 + placeholder: + en_US: Use commas to separate multiple tags + zh_Hans: 多个标签时使用半角逗号分隔 + human_description: + en_US: | + Only include tags, classes and ids from the page in the final output. Use comma separated values. Example: script, .ad, #footer + zh_Hans: | + 仅在最终输出中包含HTML页面的这些标签,可以通过标签名、类或ID来设定,使用逗号分隔值。示例:script, .ad, #footer + form: form + - name: excludeTags + type: string + label: + en_US: Exclude Tags + zh_Hans: 要移除这些标签 + human_description: + en_US: | + Tags, classes and ids to remove from the page. Use comma separated values. Example: script, .ad, #footer + zh_Hans: | + 要在最终输出中移除HTML页面的这些标签,可以通过标签名、类或ID来设定,使用逗号分隔值。示例:script, .ad, #footer + placeholder: + en_US: Use commas to separate multiple tags + zh_Hans: 多个标签时使用半角逗号分隔 + form: form + - name: onlyMainContent + type: boolean + default: false + label: + en_US: only Main Content + zh_Hans: 仅抓取主要内容 + human_description: + en_US: Only return the main content of the page excluding headers, navs, footers, etc. + zh_Hans: 只返回页面的主要内容,不包括头部、导航栏、尾部等。 + form: form + - name: waitFor + type: number + min: 0 + label: + en_US: wait For + zh_Hans: 等待时间 + human_description: + en_US: Wait x amount of milliseconds for the page to load to fetch content. + zh_Hans: 等待x毫秒以使页面加载并获取内容。 + form: form diff --git a/api/core/tools/provider/builtin/firecrawl/tools/crawl_job.py b/api/core/tools/provider/builtin/firecrawl/tools/crawl_job.py new file mode 100644 index 0000000000000000000000000000000000000000..0d2486c7ca44266b7e98de4e74b78606ec1a0801 --- /dev/null +++ b/api/core/tools/provider/builtin/firecrawl/tools/crawl_job.py @@ -0,0 +1,21 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.firecrawl.firecrawl_appx import FirecrawlApp +from core.tools.tool.builtin_tool import BuiltinTool + + +class CrawlJobTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app = FirecrawlApp( + api_key=self.runtime.credentials["firecrawl_api_key"], base_url=self.runtime.credentials["base_url"] + ) + operation = tool_parameters.get("operation", "get") + if operation == "get": + result = app.check_crawl_status(job_id=tool_parameters["job_id"]) + elif operation == "cancel": + result = app.cancel_crawl_job(job_id=tool_parameters["job_id"]) + else: + raise ValueError(f"Invalid operation: {operation}") + + return self.create_json_message(result) diff --git a/api/core/tools/provider/builtin/firecrawl/tools/crawl_job.yaml b/api/core/tools/provider/builtin/firecrawl/tools/crawl_job.yaml new file mode 100644 index 0000000000000000000000000000000000000000..78008e4ad4d8a6002d61159a34f583528b0c1344 --- /dev/null +++ b/api/core/tools/provider/builtin/firecrawl/tools/crawl_job.yaml @@ -0,0 +1,37 @@ +identity: + name: crawl_job + author: hjlarry + label: + en_US: Crawl Job + zh_Hans: 爬取任务处理 +description: + human: + en_US: Retrieve the scraping results based on the job ID, or cancel the scraping task. + zh_Hans: 根据爬取任务ID获取爬取结果,或者取消爬取任务 + llm: Retrieve the scraping results based on the job ID, or cancel the scraping task. +parameters: + - name: job_id + type: string + required: true + label: + en_US: Job ID + human_description: + en_US: Set wait_for_results to false in the Crawl tool can get the job ID. + zh_Hans: 在深度爬取工具中将等待爬取结果设置为否可以获取Job ID。 + llm_description: Set wait_for_results to false in the Crawl tool can get the job ID. + form: llm + - name: operation + type: select + required: true + options: + - value: get + label: + en_US: get crawl status + - value: cancel + label: + en_US: cancel crawl job + label: + en_US: operation + zh_Hans: 操作 + llm_description: choose the operation to perform. `get` is for getting the crawl status, `cancel` is for cancelling the crawl job. + form: llm diff --git a/api/core/tools/provider/builtin/firecrawl/tools/map.py b/api/core/tools/provider/builtin/firecrawl/tools/map.py new file mode 100644 index 0000000000000000000000000000000000000000..bdfb5faeb8e2c9727c3221fcbb9ee410b4e45984 --- /dev/null +++ b/api/core/tools/provider/builtin/firecrawl/tools/map.py @@ -0,0 +1,25 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.firecrawl.firecrawl_appx import FirecrawlApp +from core.tools.tool.builtin_tool import BuiltinTool + + +class MapTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + """ + the api doc: + https://docs.firecrawl.dev/api-reference/endpoint/map + """ + app = FirecrawlApp( + api_key=self.runtime.credentials["firecrawl_api_key"], base_url=self.runtime.credentials["base_url"] + ) + payload = {} + payload["search"] = tool_parameters.get("search") + payload["ignoreSitemap"] = tool_parameters.get("ignoreSitemap", True) + payload["includeSubdomains"] = tool_parameters.get("includeSubdomains", False) + payload["limit"] = tool_parameters.get("limit", 5000) + + map_result = app.map(url=tool_parameters["url"], **payload) + + return self.create_json_message(map_result) diff --git a/api/core/tools/provider/builtin/firecrawl/tools/map.yaml b/api/core/tools/provider/builtin/firecrawl/tools/map.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9913756983370a6c251285babb4d668cd587da5d --- /dev/null +++ b/api/core/tools/provider/builtin/firecrawl/tools/map.yaml @@ -0,0 +1,59 @@ +identity: + name: map + author: hjlarry + label: + en_US: Map + zh_Hans: 地图式快爬 +description: + human: + en_US: Input a website and get all the urls on the website - extremly fast + zh_Hans: 输入一个网站,快速获取网站上的所有网址。 + llm: Input a website and get all the urls on the website - extremly fast +parameters: + - name: url + type: string + required: true + label: + en_US: Start URL + zh_Hans: 起始URL + human_description: + en_US: The base URL to start crawling from. + zh_Hans: 要爬取网站的起始URL。 + llm_description: The URL of the website that needs to be crawled. This is a required parameter. + form: llm + - name: search + type: string + label: + en_US: search + zh_Hans: 搜索查询 + human_description: + en_US: Search query to use for mapping. During the Alpha phase, the 'smart' part of the search functionality is limited to 100 search results. However, if map finds more results, there is no limit applied. + zh_Hans: 用于映射的搜索查询。在Alpha阶段,搜索功能的“智能”部分限制为最多100个搜索结果。然而,如果地图找到了更多结果,则不施加任何限制。 + llm_description: Search query to use for mapping. During the Alpha phase, the 'smart' part of the search functionality is limited to 100 search results. However, if map finds more results, there is no limit applied. + form: llm +############## Page Options ####################### + - name: ignoreSitemap + type: boolean + default: true + label: + en_US: ignore Sitemap + zh_Hans: 忽略站点地图 + human_description: + en_US: Ignore the website sitemap when crawling. + zh_Hans: 爬取时忽略网站站点地图。 + form: form + - name: includeSubdomains + type: boolean + default: false + label: + en_US: include Subdomains + zh_Hans: 包含子域名 + form: form + - name: limit + type: number + min: 0 + default: 5000 + label: + en_US: Maximum results + zh_Hans: 最大结果数量 + form: form diff --git a/api/core/tools/provider/builtin/firecrawl/tools/scrape.py b/api/core/tools/provider/builtin/firecrawl/tools/scrape.py new file mode 100644 index 0000000000000000000000000000000000000000..f00a9b31ce8c2ccd6a7897da661931099d9e522f --- /dev/null +++ b/api/core/tools/provider/builtin/firecrawl/tools/scrape.py @@ -0,0 +1,39 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.firecrawl.firecrawl_appx import FirecrawlApp, get_array_params, get_json_params +from core.tools.tool.builtin_tool import BuiltinTool + + +class ScrapeTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> list[ToolInvokeMessage]: + """ + the api doc: + https://docs.firecrawl.dev/api-reference/endpoint/scrape + """ + app = FirecrawlApp( + api_key=self.runtime.credentials["firecrawl_api_key"], base_url=self.runtime.credentials["base_url"] + ) + + payload = {} + extract = {} + + payload["formats"] = get_array_params(tool_parameters, "formats") + payload["onlyMainContent"] = tool_parameters.get("onlyMainContent", True) + payload["includeTags"] = get_array_params(tool_parameters, "includeTags") + payload["excludeTags"] = get_array_params(tool_parameters, "excludeTags") + payload["headers"] = get_json_params(tool_parameters, "headers") + payload["waitFor"] = tool_parameters.get("waitFor", 0) + payload["timeout"] = tool_parameters.get("timeout", 30000) + + extract["schema"] = get_json_params(tool_parameters, "schema") + extract["systemPrompt"] = tool_parameters.get("systemPrompt") + extract["prompt"] = tool_parameters.get("prompt") + extract = {k: v for k, v in extract.items() if v not in (None, "")} + payload["extract"] = extract or None + + payload = {k: v for k, v in payload.items() if v not in (None, "")} + + crawl_result = app.scrape_url(url=tool_parameters["url"], **payload) + markdown_result = crawl_result.get("data", {}).get("markdown", "") + return [self.create_text_message(markdown_result), self.create_json_message(crawl_result)] diff --git a/api/core/tools/provider/builtin/firecrawl/tools/scrape.yaml b/api/core/tools/provider/builtin/firecrawl/tools/scrape.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8f1f1348a459ca2aa87d43c786008da6d3c421e3 --- /dev/null +++ b/api/core/tools/provider/builtin/firecrawl/tools/scrape.yaml @@ -0,0 +1,152 @@ +identity: + name: scrape + author: ahasasjeb + label: + en_US: Scrape + zh_Hans: 单页面抓取 +description: + human: + en_US: Turn any url into clean data. + zh_Hans: 将任何网址转换为干净的数据。 + llm: This tool is designed to scrape URL and output the content in Markdown format. +parameters: + - name: url + type: string + required: true + label: + en_US: URL to scrape + zh_Hans: 要抓取的URL + human_description: + en_US: The URL of the website to scrape and extract data from. + zh_Hans: 要抓取并提取数据的网站URL。 + llm_description: The URL of the website that needs to be crawled. This is a required parameter. + form: llm +############## Payload ####################### + - name: formats + type: string + label: + en_US: Formats + zh_Hans: 结果的格式 + placeholder: + en_US: Use commas to separate multiple tags + zh_Hans: 多个标签时使用半角逗号分隔 + human_description: + en_US: | + Formats to include in the output. Available options: markdown, html, rawHtml, links, screenshot, extract, screenshot@fullPage + zh_Hans: | + 输出中应包含的格式。可以填入: markdown, html, rawHtml, links, screenshot, extract, screenshot@fullPage + form: form + - name: onlyMainContent + type: boolean + default: false + label: + en_US: only Main Content + zh_Hans: 仅抓取主要内容 + human_description: + en_US: Only return the main content of the page excluding headers, navs, footers, etc. + zh_Hans: 只返回页面的主要内容,不包括头部、导航栏、尾部等。 + form: form + - name: includeTags + type: string + label: + en_US: Include Tags + zh_Hans: 仅抓取这些标签 + placeholder: + en_US: Use commas to separate multiple tags + zh_Hans: 多个标签时使用半角逗号分隔 + human_description: + en_US: | + Only include tags, classes and ids from the page in the final output. Use comma separated values. Example: script, .ad, #footer + zh_Hans: | + 仅在最终输出中包含HTML页面的这些标签,可以通过标签名、类或ID来设定,使用逗号分隔值。示例:script, .ad, #footer + form: form + - name: excludeTags + type: string + label: + en_US: Exclude Tags + zh_Hans: 要移除这些标签 + human_description: + en_US: | + Tags, classes and ids to remove from the page. Use comma separated values. Example: script, .ad, #footer + zh_Hans: | + 要在最终输出中移除HTML页面的这些标签,可以通过标签名、类或ID来设定,使用逗号分隔值。示例:script, .ad, #footer + placeholder: + en_US: Use commas to separate multiple tags + zh_Hans: 多个标签时使用半角逗号分隔 + form: form + - name: headers + type: string + label: + en_US: headers + zh_Hans: 请求头 + human_description: + en_US: | + Headers to send with the request. Can be used to send cookies, user-agent, etc. Example: {"cookies": "testcookies"} + zh_Hans: | + 随请求发送的头部。可以用来发送cookies、用户代理等。示例:{"cookies": "testcookies"} + placeholder: + en_US: Please enter an object that can be serialized in JSON + zh_Hans: 请输入可以json序列化的对象 + form: form + - name: waitFor + type: number + min: 0 + default: 0 + label: + en_US: wait For + zh_Hans: 等待时间 + human_description: + en_US: Wait x amount of milliseconds for the page to load to fetch content. + zh_Hans: 等待x毫秒以使页面加载并获取内容。 + form: form + - name: timeout + type: number + min: 0 + default: 30000 + label: + en_US: Timeout + human_description: + en_US: Timeout in milliseconds for the request. + zh_Hans: 请求的超时时间(以毫秒为单位)。 + form: form +############## Extractor Options ####################### + - name: schema + type: string + label: + en_US: Extractor Schema + zh_Hans: 提取时的结构 + placeholder: + en_US: Please enter an object that can be serialized in JSON + zh_Hans: 请输入可以json序列化的对象 + human_description: + en_US: | + The schema for the data to be extracted. Example: { + "type": "object", + "properties": {"company_mission": {"type": "string"}}, + "required": ["company_mission"] + } + zh_Hans: | + 使用该结构去提取,示例:{ + "type": "object", + "properties": {"company_mission": {"type": "string"}}, + "required": ["company_mission"] + } + form: form + - name: systemPrompt + type: string + label: + en_US: Extractor System Prompt + zh_Hans: 提取时的系统提示词 + human_description: + en_US: The system prompt to use for the extraction. + zh_Hans: 用于提取的系统提示。 + form: form + - name: prompt + type: string + label: + en_US: Extractor Prompt + zh_Hans: 提取时的提示词 + human_description: + en_US: The prompt to use for the extraction without a schema. + zh_Hans: 用于无schema时提取的提示词 + form: form diff --git a/api/core/tools/provider/builtin/gaode/_assets/icon.svg b/api/core/tools/provider/builtin/gaode/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..0f5729e17aea8d519f655948869fca814c97c79c --- /dev/null +++ b/api/core/tools/provider/builtin/gaode/_assets/icon.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/gaode/gaode.py b/api/core/tools/provider/builtin/gaode/gaode.py new file mode 100644 index 0000000000000000000000000000000000000000..49a8e537fb9070d73cb2db2193d44b590ae59424 --- /dev/null +++ b/api/core/tools/provider/builtin/gaode/gaode.py @@ -0,0 +1,28 @@ +import urllib.parse + +import requests + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class GaodeProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + try: + if "api_key" not in credentials or not credentials.get("api_key"): + raise ToolProviderCredentialValidationError("Gaode API key is required.") + + try: + response = requests.get( + url="https://restapi.amap.com/v3/geocode/geo?address={address}&key={apikey}".format( + address=urllib.parse.quote("广东省广州市天河区广州塔"), apikey=credentials.get("api_key") + ) + ) + if response.status_code == 200 and (response.json()).get("info") == "OK": + pass + else: + raise ToolProviderCredentialValidationError((response.json()).get("info")) + except Exception as e: + raise ToolProviderCredentialValidationError("Gaode API Key is invalid. {}".format(e)) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/gaode/gaode.yaml b/api/core/tools/provider/builtin/gaode/gaode.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2eb3b161a29915e4ee8a0b1626eace636dd7792c --- /dev/null +++ b/api/core/tools/provider/builtin/gaode/gaode.yaml @@ -0,0 +1,34 @@ +identity: + author: CharlieWei + name: gaode + label: + en_US: Autonavi + zh_Hans: 高德 + pt_BR: Autonavi + description: + en_US: Autonavi Open Platform service toolkit. + zh_Hans: 高德开放平台服务工具包。 + pt_BR: Kit de ferramentas de serviço Autonavi Open Platform. + icon: icon.svg + tags: + - utilities + - productivity + - travel + - weather +credentials_for_provider: + api_key: + type: secret-input + required: true + label: + en_US: API Key + zh_Hans: API Key + pt_BR: Fogo a chave + placeholder: + en_US: Please enter your Autonavi API Key + zh_Hans: 请输入你的高德开放平台 API Key + pt_BR: Insira sua chave de API Autonavi + help: + en_US: Get your API Key from Autonavi + zh_Hans: 从高德获取您的 API Key + pt_BR: Obtenha sua chave de API do Autonavi + url: https://console.amap.com/dev/key/app diff --git a/api/core/tools/provider/builtin/gaode/tools/gaode_weather.py b/api/core/tools/provider/builtin/gaode/tools/gaode_weather.py new file mode 100644 index 0000000000000000000000000000000000000000..4642415e6dd394a3112a6aaa913d9e50f41ea38f --- /dev/null +++ b/api/core/tools/provider/builtin/gaode/tools/gaode_weather.py @@ -0,0 +1,65 @@ +import json +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class GaodeRepositoriesTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + city = tool_parameters.get("city", "") + if not city: + return self.create_text_message("Please tell me your city") + + if "api_key" not in self.runtime.credentials or not self.runtime.credentials.get("api_key"): + return self.create_text_message("Gaode API key is required.") + + try: + s = requests.session() + api_domain = "https://restapi.amap.com/v3" + city_response = s.request( + method="GET", + headers={"Content-Type": "application/json; charset=utf-8"}, + url="{url}/config/district?keywords={keywords}&subdistrict=0&extensions=base&key={apikey}".format( + url=api_domain, keywords=city, apikey=self.runtime.credentials.get("api_key") + ), + ) + City_data = city_response.json() + if city_response.status_code == 200 and City_data.get("info") == "OK": + if len(City_data.get("districts")) > 0: + CityCode = City_data["districts"][0]["adcode"] + weatherInfo_response = s.request( + method="GET", + url="{url}/weather/weatherInfo?city={citycode}&extensions=all&key={apikey}&output=json".format( + url=api_domain, citycode=CityCode, apikey=self.runtime.credentials.get("api_key") + ), + ) + weatherInfo_data = weatherInfo_response.json() + if weatherInfo_response.status_code == 200 and weatherInfo_data.get("info") == "OK": + contents = [] + if len(weatherInfo_data.get("forecasts")) > 0: + for item in weatherInfo_data["forecasts"][0]["casts"]: + content = {} + content["date"] = item.get("date") + content["week"] = item.get("week") + content["dayweather"] = item.get("dayweather") + content["daytemp_float"] = item.get("daytemp_float") + content["daywind"] = item.get("daywind") + content["nightweather"] = item.get("nightweather") + content["nighttemp_float"] = item.get("nighttemp_float") + contents.append(content) + s.close() + return self.create_text_message( + self.summary(user_id=user_id, content=json.dumps(contents, ensure_ascii=False)) + ) + s.close() + return self.create_text_message(f"No weather information for {city} was found.") + except Exception as e: + return self.create_text_message("Gaode API Key and Api Version is invalid. {}".format(e)) diff --git a/api/core/tools/provider/builtin/gaode/tools/gaode_weather.yaml b/api/core/tools/provider/builtin/gaode/tools/gaode_weather.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e41851e188edeecb2787fa5b14c693e16018ed42 --- /dev/null +++ b/api/core/tools/provider/builtin/gaode/tools/gaode_weather.yaml @@ -0,0 +1,28 @@ +identity: + name: gaode_weather + author: CharlieWei + label: + en_US: Weather Forecast + zh_Hans: 天气预报 + pt_BR: Previsão do tempo + icon: icon.svg +description: + human: + en_US: Weather forecast inquiry + zh_Hans: 天气预报查询。 + pt_BR: Inquérito sobre previsão meteorológica. + llm: A tool when you want to ask about the weather or weather-related question. +parameters: + - name: city + type: string + required: true + label: + en_US: city + zh_Hans: 城市 + pt_BR: cidade + human_description: + en_US: Target city for weather forecast query. + zh_Hans: 天气预报查询的目标城市。 + pt_BR: Cidade de destino para consulta de previsão do tempo. + llm_description: If you don't know you can extract the city name from the question or you can reply:Please tell me your city. You have to extract the Chinese city name from the question. + form: llm diff --git a/api/core/tools/provider/builtin/getimgai/_assets/icon.svg b/api/core/tools/provider/builtin/getimgai/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..6b2513386da458116336e0ca41a7c7353e696956 --- /dev/null +++ b/api/core/tools/provider/builtin/getimgai/_assets/icon.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/getimgai/getimgai.py b/api/core/tools/provider/builtin/getimgai/getimgai.py new file mode 100644 index 0000000000000000000000000000000000000000..bbd07d120fd0ea75c28ebde2266c68e1ff723cb5 --- /dev/null +++ b/api/core/tools/provider/builtin/getimgai/getimgai.py @@ -0,0 +1,19 @@ +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.getimgai.tools.text2image import Text2ImageTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class GetImgAIProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + try: + # Example validation using the text2image tool + Text2ImageTool().fork_tool_runtime(runtime={"credentials": credentials}).invoke( + user_id="", + tool_parameters={ + "prompt": "A fire egg", + "response_format": "url", + "style": "photorealism", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/getimgai/getimgai.yaml b/api/core/tools/provider/builtin/getimgai/getimgai.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c9db0a9e22a6c467833caabe7780b85c23d6e3bb --- /dev/null +++ b/api/core/tools/provider/builtin/getimgai/getimgai.yaml @@ -0,0 +1,29 @@ +identity: + author: Matri Qi + name: getimgai + label: + en_US: getimg.ai + zh_CN: getimg.ai + description: + en_US: GetImg API integration for image generation and scraping. + icon: icon.svg + tags: + - image +credentials_for_provider: + getimg_api_key: + type: secret-input + required: true + label: + en_US: getimg.ai API Key + placeholder: + en_US: Please input your getimg.ai API key + help: + en_US: Get your getimg.ai API key from your getimg.ai account settings. If you are using a self-hosted version, you may enter any key at your convenience. + url: https://dashboard.getimg.ai/api-keys + base_url: + type: text-input + required: false + label: + en_US: getimg.ai server's Base URL + placeholder: + en_US: https://api.getimg.ai/v1 diff --git a/api/core/tools/provider/builtin/getimgai/getimgai_appx.py b/api/core/tools/provider/builtin/getimgai/getimgai_appx.py new file mode 100644 index 0000000000000000000000000000000000000000..0e95a5f654505f5023ed3dfd600df196d9153345 --- /dev/null +++ b/api/core/tools/provider/builtin/getimgai/getimgai_appx.py @@ -0,0 +1,55 @@ +import logging +import time +from collections.abc import Mapping +from typing import Any + +import requests +from requests.exceptions import HTTPError + +logger = logging.getLogger(__name__) + + +class GetImgAIApp: + def __init__(self, api_key: str | None = None, base_url: str | None = None): + self.api_key = api_key + self.base_url = base_url or "https://api.getimg.ai/v1" + if not self.api_key: + raise ValueError("API key is required") + + def _prepare_headers(self): + headers = {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"} + return headers + + def _request( + self, + method: str, + url: str, + data: Mapping[str, Any] | None = None, + headers: Mapping[str, str] | None = None, + retries: int = 3, + backoff_factor: float = 0.3, + ) -> Mapping[str, Any] | None: + for i in range(retries): + try: + response = requests.request(method, url, json=data, headers=headers) + response.raise_for_status() + return response.json() + except requests.exceptions.RequestException as e: + if i < retries - 1 and isinstance(e, HTTPError) and e.response.status_code >= 500: + time.sleep(backoff_factor * (2**i)) + else: + raise + return None + + def text2image(self, mode: str, **kwargs): + data = kwargs["params"] + if not data.get("prompt"): + raise ValueError("Prompt is required") + + endpoint = f"{self.base_url}/{mode}/text-to-image" + headers = self._prepare_headers() + logger.debug(f"Send request to {endpoint=} body={data}") + response = self._request("POST", endpoint, data, headers) + if response is None: + raise HTTPError("Failed to initiate getimg.ai after multiple retries") + return response diff --git a/api/core/tools/provider/builtin/getimgai/tools/text2image.py b/api/core/tools/provider/builtin/getimgai/tools/text2image.py new file mode 100644 index 0000000000000000000000000000000000000000..c556749552c8ef0df267a6af440556516e8566fe --- /dev/null +++ b/api/core/tools/provider/builtin/getimgai/tools/text2image.py @@ -0,0 +1,39 @@ +import json +from typing import Any, Union + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.getimgai.getimgai_appx import GetImgAIApp +from core.tools.tool.builtin_tool import BuiltinTool + + +class Text2ImageTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + app = GetImgAIApp( + api_key=self.runtime.credentials["getimg_api_key"], base_url=self.runtime.credentials["base_url"] + ) + + options = { + "style": tool_parameters.get("style"), + "prompt": tool_parameters.get("prompt"), + "aspect_ratio": tool_parameters.get("aspect_ratio"), + "output_format": tool_parameters.get("output_format", "jpeg"), + "response_format": tool_parameters.get("response_format", "url"), + "width": tool_parameters.get("width"), + "height": tool_parameters.get("height"), + "steps": tool_parameters.get("steps"), + "negative_prompt": tool_parameters.get("negative_prompt"), + "prompt_2": tool_parameters.get("prompt_2"), + } + options = {k: v for k, v in options.items() if v} + + text2image_result = app.text2image(mode=tool_parameters.get("mode", "essential-v2"), params=options, wait=True) + + if not isinstance(text2image_result, str): + text2image_result = json.dumps(text2image_result, ensure_ascii=False, indent=4) + + if not text2image_result: + return self.create_text_message("getimg.ai request failed.") + + return self.create_text_message(text2image_result) diff --git a/api/core/tools/provider/builtin/getimgai/tools/text2image.yaml b/api/core/tools/provider/builtin/getimgai/tools/text2image.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d972186f56d6a6d7f960b48199fced4d86a0d660 --- /dev/null +++ b/api/core/tools/provider/builtin/getimgai/tools/text2image.yaml @@ -0,0 +1,167 @@ +identity: + name: text2image + author: Matri Qi + label: + en_US: text2image + icon: icon.svg +description: + human: + en_US: Generate image via getimg.ai. + llm: This tool is used to generate image from prompt or image via https://getimg.ai. +parameters: + - name: prompt + type: string + required: true + label: + en_US: prompt + human_description: + en_US: The text prompt used to generate the image. The getimg.aier will generate an image based on this prompt. + llm_description: this prompt text will be used to generate image. + form: llm + - name: mode + type: select + required: false + label: + en_US: mode + human_description: + en_US: The getimg.ai mode to use. The mode determines the endpoint used to generate the image. + form: form + options: + - value: "essential-v2" + label: + en_US: essential-v2 + - value: stable-diffusion-xl + label: + en_US: stable-diffusion-xl + - value: stable-diffusion + label: + en_US: stable-diffusion + - value: latent-consistency + label: + en_US: latent-consistency + - name: style + type: select + required: false + label: + en_US: style + human_description: + en_US: The style preset to use. The style preset guides the generation towards a particular style. It's just efficient for `Essential V2` mode. + form: form + options: + - value: photorealism + label: + en_US: photorealism + - value: anime + label: + en_US: anime + - value: art + label: + en_US: art + - name: aspect_ratio + type: select + required: false + label: + en_US: "aspect ratio" + human_description: + en_US: The aspect ratio of the generated image. It's just efficient for `Essential V2` mode. + form: form + options: + - value: "1:1" + label: + en_US: "1:1" + - value: "4:5" + label: + en_US: "4:5" + - value: "5:4" + label: + en_US: "5:4" + - value: "2:3" + label: + en_US: "2:3" + - value: "3:2" + label: + en_US: "3:2" + - value: "4:7" + label: + en_US: "4:7" + - value: "7:4" + label: + en_US: "7:4" + - name: output_format + type: select + required: false + label: + en_US: "output format" + human_description: + en_US: The file format of the generated image. + form: form + options: + - value: jpeg + label: + en_US: jpeg + - value: png + label: + en_US: png + - name: response_format + type: select + required: false + label: + en_US: "response format" + human_description: + en_US: The format in which the generated images are returned. Must be one of url or b64. URLs are only valid for 1 hour after the image has been generated. + form: form + options: + - value: url + label: + en_US: url + - value: b64 + label: + en_US: b64 + - name: model + type: string + required: false + label: + en_US: model + human_description: + en_US: Model ID supported by this pipeline and family. It's just efficient for `Stable Diffusion XL`, `Stable Diffusion`, `Latent Consistency` mode. + form: form + - name: negative_prompt + type: string + required: false + label: + en_US: negative prompt + human_description: + en_US: Text input that will not guide the image generation. It's just efficient for `Stable Diffusion XL`, `Stable Diffusion`, `Latent Consistency` mode. + form: form + - name: prompt_2 + type: string + required: false + label: + en_US: prompt2 + human_description: + en_US: Prompt sent to second tokenizer and text encoder. If not defined, prompt is used in both text-encoders. It's just efficient for `Stable Diffusion XL` mode. + form: form + - name: width + type: number + required: false + label: + en_US: width + human_description: + en_US: he width of the generated image in pixels. Width needs to be multiple of 64. + form: form + - name: height + type: number + required: false + label: + en_US: height + human_description: + en_US: he height of the generated image in pixels. Height needs to be multiple of 64. + form: form + - name: steps + type: number + required: false + label: + en_US: steps + human_description: + en_US: The number of denoising steps. More steps usually can produce higher quality images, but take more time to generate. It's just efficient for `Stable Diffusion XL`, `Stable Diffusion`, `Latent Consistency` mode. + form: form diff --git a/api/core/tools/provider/builtin/gitee_ai/_assets/icon.svg b/api/core/tools/provider/builtin/gitee_ai/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..6dd75d1a6b5b447b0225ee71b23bef128c918c38 --- /dev/null +++ b/api/core/tools/provider/builtin/gitee_ai/_assets/icon.svg @@ -0,0 +1,3 @@ + + + diff --git a/api/core/tools/provider/builtin/gitee_ai/gitee_ai.py b/api/core/tools/provider/builtin/gitee_ai/gitee_ai.py new file mode 100644 index 0000000000000000000000000000000000000000..151cafec14b2b7b64e55ece84e766c65f2ac00d3 --- /dev/null +++ b/api/core/tools/provider/builtin/gitee_ai/gitee_ai.py @@ -0,0 +1,17 @@ +import requests + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class GiteeAIProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + url = "https://ai.gitee.com/api/base/account/me" + headers = { + "accept": "application/json", + "authorization": f"Bearer {credentials.get('api_key')}", + } + + response = requests.get(url, headers=headers) + if response.status_code != 200: + raise ToolProviderCredentialValidationError("GiteeAI API key is invalid") diff --git a/api/core/tools/provider/builtin/gitee_ai/gitee_ai.yaml b/api/core/tools/provider/builtin/gitee_ai/gitee_ai.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d0475665dd7ac7838c934c71c6c45319644d763f --- /dev/null +++ b/api/core/tools/provider/builtin/gitee_ai/gitee_ai.yaml @@ -0,0 +1,22 @@ +identity: + author: Gitee AI + name: gitee_ai + label: + en_US: Gitee AI + zh_Hans: Gitee AI + description: + en_US: Quickly experience large models and explore the leading AI open source world + zh_Hans: 快速体验大模型,领先探索 AI 开源世界 + icon: icon.svg + tags: + - image +credentials_for_provider: + api_key: + type: secret-input + required: true + label: + en_US: API Key + placeholder: + zh_Hans: 在此输入您的 API Key + en_US: Enter your API Key + url: https://ai.gitee.com/dashboard/settings/tokens diff --git a/api/core/tools/provider/builtin/gitee_ai/tools/embedding.py b/api/core/tools/provider/builtin/gitee_ai/tools/embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..ab03759c1966ad4cf8d5ca6dad7397578f37ff11 --- /dev/null +++ b/api/core/tools/provider/builtin/gitee_ai/tools/embedding.py @@ -0,0 +1,25 @@ +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class GiteeAIToolEmbedding(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + headers = { + "content-type": "application/json", + "authorization": f"Bearer {self.runtime.credentials['api_key']}", + } + + payload = {"inputs": tool_parameters.get("inputs")} + model = tool_parameters.get("model", "bge-m3") + url = f"https://ai.gitee.com/api/serverless/{model}/embeddings" + response = requests.post(url, json=payload, headers=headers) + if response.status_code != 200: + return self.create_text_message(f"Got Error Response:{response.text}") + + return [self.create_text_message(response.content.decode("utf-8"))] diff --git a/api/core/tools/provider/builtin/gitee_ai/tools/embedding.yaml b/api/core/tools/provider/builtin/gitee_ai/tools/embedding.yaml new file mode 100644 index 0000000000000000000000000000000000000000..53e569d731d072d44c1b6593034a76c93adf146e --- /dev/null +++ b/api/core/tools/provider/builtin/gitee_ai/tools/embedding.yaml @@ -0,0 +1,37 @@ +identity: + name: embedding + author: gitee_ai + label: + en_US: embedding + icon: icon.svg +description: + human: + en_US: Generate word embeddings using Serverless-supported models (compatible with OpenAI) + llm: This tool is used to generate word embeddings from text input. +parameters: + - name: model + type: string + required: true + in: path + description: + en_US: Supported Embedding (compatible with OpenAI) interface models + enum: + - bge-m3 + - bge-large-zh-v1.5 + - bge-small-zh-v1.5 + label: + en_US: Service Model + zh_Hans: 服务模型 + default: bge-m3 + form: form + - name: inputs + type: string + required: true + label: + en_US: Input Text + zh_Hans: 输入文本 + human_description: + en_US: The text input used to generate embeddings. + zh_Hans: 用于生成词向量的输入文本。 + llm_description: This text input will be used to generate embeddings. + form: llm diff --git a/api/core/tools/provider/builtin/gitee_ai/tools/risk-control.py b/api/core/tools/provider/builtin/gitee_ai/tools/risk-control.py new file mode 100644 index 0000000000000000000000000000000000000000..e3558ce69915be3ca95c7f08150b7b4fb733ccb6 --- /dev/null +++ b/api/core/tools/provider/builtin/gitee_ai/tools/risk-control.py @@ -0,0 +1,26 @@ +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class GiteeAIToolRiskControl(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + headers = { + "content-type": "application/json", + "authorization": f"Bearer {self.runtime.credentials['api_key']}", + } + + inputs = [{"type": "text", "text": tool_parameters.get("input-text")}] + model = tool_parameters.get("model", "Security-semantic-filtering") + payload = {"model": model, "input": inputs} + url = "https://ai.gitee.com/v1/moderations" + response = requests.post(url, json=payload, headers=headers) + if response.status_code != 200: + return self.create_text_message(f"Got Error Response:{response.text}") + + return [self.create_text_message(response.content.decode("utf-8"))] diff --git a/api/core/tools/provider/builtin/gitee_ai/tools/risk-control.yaml b/api/core/tools/provider/builtin/gitee_ai/tools/risk-control.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6e7229dc1c54d5b59355adfb026d05d2d635382f --- /dev/null +++ b/api/core/tools/provider/builtin/gitee_ai/tools/risk-control.yaml @@ -0,0 +1,32 @@ +identity: + name: risk control + author: gitee_ai + label: + en_US: risk control identification + zh_Hans: 风控识别 + icon: icon.svg +description: + human: + en_US: Ensuring the protection and compliance of sensitive information through the filtering and analysis of data semantics + zh_Hans: 通过对数据语义的过滤和分析,确保敏感信息的保护和合规性 + llm: This tool is used to risk control identification. +parameters: + - name: model + type: string + required: true + default: Security-semantic-filtering + label: + en_US: Service Model + zh_Hans: 服务模型 + form: form + - name: input-text + type: string + required: true + label: + en_US: Input Text + zh_Hans: 输入文本 + human_description: + en_US: The text input for filtering and analysis. + zh_Hans: 用于分析过滤的文本 + llm_description: The text input for filtering and analysis. + form: llm diff --git a/api/core/tools/provider/builtin/gitee_ai/tools/text-to-image.py b/api/core/tools/provider/builtin/gitee_ai/tools/text-to-image.py new file mode 100644 index 0000000000000000000000000000000000000000..bb0b2c915b333ccb2a79e36574ce33ccc9285444 --- /dev/null +++ b/api/core/tools/provider/builtin/gitee_ai/tools/text-to-image.py @@ -0,0 +1,33 @@ +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class GiteeAIToolText2Image(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + headers = { + "content-type": "application/json", + "authorization": f"Bearer {self.runtime.credentials['api_key']}", + } + + payload = { + "inputs": tool_parameters.get("inputs"), + "width": tool_parameters.get("width", "720"), + "height": tool_parameters.get("height", "720"), + } + model = tool_parameters.get("model", "Kolors") + url = f"https://ai.gitee.com/api/serverless/{model}/text-to-image" + + response = requests.post(url, json=payload, headers=headers) + if response.status_code != 200: + return self.create_text_message(f"Got Error Response:{response.text}") + + # The returned image is base64 and needs to be mark as an image + result = [self.create_blob_message(blob=response.content, meta={"mime_type": "image/jpeg"})] + + return result diff --git a/api/core/tools/provider/builtin/gitee_ai/tools/text-to-image.yaml b/api/core/tools/provider/builtin/gitee_ai/tools/text-to-image.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5e03f9abe9dfe4e242aef7d0d15ac5a99efa55b3 --- /dev/null +++ b/api/core/tools/provider/builtin/gitee_ai/tools/text-to-image.yaml @@ -0,0 +1,72 @@ +identity: + name: text to image + author: gitee_ai + label: + en_US: text to image + icon: icon.svg +description: + human: + en_US: generate images using a variety of popular models + llm: This tool is used to generate image from text. +parameters: + - name: model + type: select + required: true + options: + - value: flux-1-schnell + label: + en_US: flux-1-schnell + - value: Kolors + label: + en_US: Kolors + - value: stable-diffusion-3-medium + label: + en_US: stable-diffusion-3-medium + - value: stable-diffusion-xl-base-1.0 + label: + en_US: stable-diffusion-xl-base-1.0 + - value: stable-diffusion-v1-4 + label: + en_US: stable-diffusion-v1-4 + default: Kolors + label: + en_US: Choose Image Model + zh_Hans: 选择生成图片的模型 + form: form + - name: inputs + type: string + required: true + label: + en_US: Input Text + zh_Hans: 输入文本 + human_description: + en_US: The text input used to generate the image. + zh_Hans: 用于生成图片的输入文本。 + llm_description: This text input will be used to generate image. + form: llm + - name: width + type: number + required: true + default: 720 + min: 1 + max: 1024 + label: + en_US: Image Width + zh_Hans: 图片宽度 + human_description: + en_US: The width of the generated image. + zh_Hans: 生成图片的宽度。 + form: form + - name: height + type: number + required: true + default: 720 + min: 1 + max: 1024 + label: + en_US: Image Height + zh_Hans: 图片高度 + human_description: + en_US: The height of the generated image. + zh_Hans: 生成图片的高度。 + form: form diff --git a/api/core/tools/provider/builtin/github/_assets/icon.svg b/api/core/tools/provider/builtin/github/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..d56adb2c2f9955d1d22e82871775d925f97d6403 --- /dev/null +++ b/api/core/tools/provider/builtin/github/_assets/icon.svg @@ -0,0 +1,17 @@ + + + github [#142] + Created with Sketch. + + + + + + + + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/github/github.py b/api/core/tools/provider/builtin/github/github.py new file mode 100644 index 0000000000000000000000000000000000000000..87a34ac3e806ea9d3cd26ca45180c8035a278e47 --- /dev/null +++ b/api/core/tools/provider/builtin/github/github.py @@ -0,0 +1,32 @@ +import requests + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class GithubProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + try: + if "access_tokens" not in credentials or not credentials.get("access_tokens"): + raise ToolProviderCredentialValidationError("Github API Access Tokens is required.") + if "api_version" not in credentials or not credentials.get("api_version"): + api_version = "2022-11-28" + else: + api_version = credentials.get("api_version") + + try: + headers = { + "Content-Type": "application/vnd.github+json", + "Authorization": f"Bearer {credentials.get('access_tokens')}", + "X-GitHub-Api-Version": api_version, + } + + response = requests.get( + url="https://api.github.com/search/users?q={account}".format(account="charli117"), headers=headers + ) + if response.status_code != 200: + raise ToolProviderCredentialValidationError((response.json()).get("message")) + except Exception as e: + raise ToolProviderCredentialValidationError("Github API Key and Api Version is invalid. {}".format(e)) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/github/github.yaml b/api/core/tools/provider/builtin/github/github.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c3d85fc3f69cf72bca18b50e365cfad29e69459f --- /dev/null +++ b/api/core/tools/provider/builtin/github/github.yaml @@ -0,0 +1,48 @@ +identity: + author: CharlieWei + name: github + label: + en_US: Github + zh_Hans: Github + pt_BR: Github + description: + en_US: GitHub is an online software source code hosting service. + zh_Hans: GitHub是一个在线软件源代码托管服务平台。 + pt_BR: GitHub é uma plataforma online para serviços de hospedagem de código fonte de software. + icon: icon.svg + tags: + - utilities +credentials_for_provider: + access_tokens: + type: secret-input + required: true + label: + en_US: Access Tokens + zh_Hans: Access Tokens + pt_BR: Tokens de acesso + placeholder: + en_US: Please input your Github Access Tokens + zh_Hans: 请输入你的 Github Access Tokens + pt_BR: Insira seus Tokens de Acesso do Github + help: + en_US: Get your Access Tokens from Github + zh_Hans: 从 Github 获取您的 Access Tokens + pt_BR: Obtenha sua chave da API do Google no Google + url: https://github.com/settings/tokens?type=beta + api_version: + type: text-input + required: false + default: '2022-11-28' + label: + en_US: API Version + zh_Hans: API Version + pt_BR: Versão da API + placeholder: + en_US: Please input your Github API Version + zh_Hans: 请输入你的 Github API Version + pt_BR: Insira sua versão da API do Github + help: + en_US: Get your API Version from Github + zh_Hans: 从 Github 获取您的 API Version + pt_BR: Obtenha sua versão da API do Github + url: https://docs.github.com/en/rest/about-the-rest-api/api-versions?apiVersion=2022-11-28 diff --git a/api/core/tools/provider/builtin/github/tools/github_repositories.py b/api/core/tools/provider/builtin/github/tools/github_repositories.py new file mode 100644 index 0000000000000000000000000000000000000000..32f9922e651785ccf288dd7889297f4ecb287e37 --- /dev/null +++ b/api/core/tools/provider/builtin/github/tools/github_repositories.py @@ -0,0 +1,70 @@ +import json +from datetime import datetime +from typing import Any, Union +from urllib.parse import quote + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class GithubRepositoriesTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + top_n = tool_parameters.get("top_n", 5) + query = tool_parameters.get("query", "") + if not query: + return self.create_text_message("Please input symbol") + + if "access_tokens" not in self.runtime.credentials or not self.runtime.credentials.get("access_tokens"): + return self.create_text_message("Github API Access Tokens is required.") + if "api_version" not in self.runtime.credentials or not self.runtime.credentials.get("api_version"): + api_version = "2022-11-28" + else: + api_version = self.runtime.credentials.get("api_version") + + try: + headers = { + "Content-Type": "application/vnd.github+json", + "Authorization": f"Bearer {self.runtime.credentials.get('access_tokens')}", + "X-GitHub-Api-Version": api_version, + } + s = requests.session() + api_domain = "https://api.github.com" + response = s.request( + method="GET", + headers=headers, + url=f"{api_domain}/search/repositories?q={quote(query)}&sort=stars&per_page={top_n}&order=desc", + ) + response_data = response.json() + if response.status_code == 200 and isinstance(response_data.get("items"), list): + contents = [] + if len(response_data.get("items")) > 0: + for item in response_data.get("items"): + content = {} + updated_at_object = datetime.strptime(item["updated_at"], "%Y-%m-%dT%H:%M:%SZ") + content["owner"] = item["owner"]["login"] + content["name"] = item["name"] + content["description"] = ( + item["description"][:100] + "..." if len(item["description"]) > 100 else item["description"] + ) + content["url"] = item["html_url"] + content["star"] = item["watchers"] + content["forks"] = item["forks"] + content["updated"] = updated_at_object.strftime("%Y-%m-%d") + contents.append(content) + s.close() + return self.create_text_message( + self.summary(user_id=user_id, content=json.dumps(contents, ensure_ascii=False)) + ) + else: + return self.create_text_message(f"No items related to {query} were found.") + else: + return self.create_text_message((response.json()).get("message")) + except Exception as e: + return self.create_text_message("Github API Key and Api Version is invalid. {}".format(e)) diff --git a/api/core/tools/provider/builtin/github/tools/github_repositories.yaml b/api/core/tools/provider/builtin/github/tools/github_repositories.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c170aee797fe4df8959db541271b8a97a9cc5b3c --- /dev/null +++ b/api/core/tools/provider/builtin/github/tools/github_repositories.yaml @@ -0,0 +1,42 @@ +identity: + name: github_repositories + author: CharlieWei + label: + en_US: Search Repositories + zh_Hans: 仓库搜索 + pt_BR: Pesquisar Repositórios + icon: icon.svg +description: + human: + en_US: Search the Github repository to retrieve the open source projects you need + zh_Hans: 搜索Github仓库,检索你需要的开源项目。 + pt_BR: Pesquise o repositório do Github para recuperar os projetos de código aberto necessários. + llm: A tool when you wants to search for popular warehouses or open source projects for any keyword. format query condition like "keywords+language:js", language can be other dev languages. +parameters: + - name: query + type: string + required: true + label: + en_US: query + zh_Hans: 关键字 + pt_BR: consulta + human_description: + en_US: You want to find the project development language, keywords, For example. Find 10 Python developed PDF document parsing projects. + zh_Hans: 你想要找的项目开发语言、关键字,如:找10个Python开发的PDF文档解析项目。 + pt_BR: Você deseja encontrar a linguagem de desenvolvimento do projeto, palavras-chave, Por exemplo. Encontre 10 projetos de análise de documentos PDF desenvolvidos em Python. + llm_description: The query of you want to search, format query condition like "keywords+language:js", language can be other dev languages. + form: llm + - name: top_n + type: number + default: 5 + required: true + label: + en_US: Top N + zh_Hans: Top N + pt_BR: Topo N + human_description: + en_US: Number of records returned by sorting based on stars. 5 is returned by default. + zh_Hans: 基于stars排序返回的记录数, 默认返回5条。 + pt_BR: Número de registros retornados por classificação com base em estrelas. 5 é retornado por padrão. + llm_description: Extract the first N records from the returned result. + form: llm diff --git a/api/core/tools/provider/builtin/gitlab/_assets/gitlab.svg b/api/core/tools/provider/builtin/gitlab/_assets/gitlab.svg new file mode 100644 index 0000000000000000000000000000000000000000..07734077d5d300fe90bcfa8067fd7214a89ffc52 --- /dev/null +++ b/api/core/tools/provider/builtin/gitlab/_assets/gitlab.svg @@ -0,0 +1,2 @@ + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/gitlab/gitlab.py b/api/core/tools/provider/builtin/gitlab/gitlab.py new file mode 100644 index 0000000000000000000000000000000000000000..9bd4a0bd52ea6452ec7ae03713d13908cbca21f1 --- /dev/null +++ b/api/core/tools/provider/builtin/gitlab/gitlab.py @@ -0,0 +1,32 @@ +from typing import Any + +import requests + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class GitlabProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + if "access_tokens" not in credentials or not credentials.get("access_tokens"): + raise ToolProviderCredentialValidationError("Gitlab Access Tokens is required.") + + if "site_url" not in credentials or not credentials.get("site_url"): + site_url = "https://gitlab.com" + else: + site_url = credentials.get("site_url") + + try: + headers = { + "Content-Type": "application/vnd.text+json", + "Authorization": f"Bearer {credentials.get('access_tokens')}", + } + + response = requests.get(url=f"{site_url}/api/v4/user", headers=headers) + if response.status_code != 200: + raise ToolProviderCredentialValidationError((response.json()).get("message")) + except Exception as e: + raise ToolProviderCredentialValidationError("Gitlab Access Tokens is invalid. {}".format(e)) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/gitlab/gitlab.yaml b/api/core/tools/provider/builtin/gitlab/gitlab.yaml new file mode 100644 index 0000000000000000000000000000000000000000..22d7ebf73ac2aa4a723bc3e6d8c5d1b81bf2004c --- /dev/null +++ b/api/core/tools/provider/builtin/gitlab/gitlab.yaml @@ -0,0 +1,38 @@ +identity: + author: Leo.Wang + name: gitlab + label: + en_US: GitLab + zh_Hans: GitLab + description: + en_US: GitLab plugin, API v4 only. + zh_Hans: 用于获取GitLab内容的插件,目前仅支持 API v4。 + icon: gitlab.svg +credentials_for_provider: + access_tokens: + type: secret-input + required: true + label: + en_US: GitLab access token + zh_Hans: GitLab access token + placeholder: + en_US: Please input your GitLab access token + zh_Hans: 请输入你的 GitLab access token + help: + en_US: Get your GitLab access token from GitLab + zh_Hans: 从 GitLab 获取您的 access token + url: https://docs.gitlab.com/16.9/ee/api/oauth2.html + site_url: + type: text-input + required: false + default: 'https://gitlab.com' + label: + en_US: GitLab site url + zh_Hans: GitLab site url + placeholder: + en_US: Please input your GitLab site url + zh_Hans: 请输入你的 GitLab site url + help: + en_US: Find your GitLab url + zh_Hans: 找到你的 GitLab url + url: https://gitlab.com/help diff --git a/api/core/tools/provider/builtin/gitlab/tools/gitlab_commits.py b/api/core/tools/provider/builtin/gitlab/tools/gitlab_commits.py new file mode 100644 index 0000000000000000000000000000000000000000..716da7c8c110f38873047e2782b453d3e8497976 --- /dev/null +++ b/api/core/tools/provider/builtin/gitlab/tools/gitlab_commits.py @@ -0,0 +1,134 @@ +import json +import urllib.parse +from datetime import datetime, timedelta +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class GitlabCommitsTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + branch = tool_parameters.get("branch", "") + repository = tool_parameters.get("repository", "") + employee = tool_parameters.get("employee", "") + start_time = tool_parameters.get("start_time", "") + end_time = tool_parameters.get("end_time", "") + change_type = tool_parameters.get("change_type", "all") + + if not repository: + return self.create_text_message("Either repository is required") + + if not start_time: + start_time = (datetime.utcnow() - timedelta(days=1)).isoformat() + if not end_time: + end_time = datetime.utcnow().isoformat() + + access_token = self.runtime.credentials.get("access_tokens") + site_url = self.runtime.credentials.get("site_url") + + if "access_tokens" not in self.runtime.credentials or not self.runtime.credentials.get("access_tokens"): + return self.create_text_message("Gitlab API Access Tokens is required.") + if "site_url" not in self.runtime.credentials or not self.runtime.credentials.get("site_url"): + site_url = "https://gitlab.com" + + # Get commit content + result = self.fetch_commits( + site_url, access_token, repository, branch, employee, start_time, end_time, change_type, is_repository=True + ) + + return [self.create_json_message(item) for item in result] + + def fetch_commits( + self, + site_url: str, + access_token: str, + repository: str, + branch: str, + employee: str, + start_time: str, + end_time: str, + change_type: str, + is_repository: bool, + ) -> list[dict[str, Any]]: + domain = site_url + headers = {"PRIVATE-TOKEN": access_token} + results = [] + + try: + # URL encode the repository path + encoded_repository = urllib.parse.quote(repository, safe="") + commits_url = f"{domain}/api/v4/projects/{encoded_repository}/repository/commits" + + # Fetch commits for the repository + params = {"since": start_time, "until": end_time} + if branch: + params["ref_name"] = branch + if employee: + params["author"] = employee + + commits_response = requests.get(commits_url, headers=headers, params=params) + commits_response.raise_for_status() + commits = commits_response.json() + + for commit in commits: + commit_sha = commit["id"] + author_name = commit["author_name"] + + diff_url = f"{domain}/api/v4/projects/{encoded_repository}/repository/commits/{commit_sha}/diff" + + diff_response = requests.get(diff_url, headers=headers) + diff_response.raise_for_status() + diffs = diff_response.json() + + for diff in diffs: + # Calculate code lines of changes + added_lines = diff["diff"].count("\n+") + removed_lines = diff["diff"].count("\n-") + total_changes = added_lines + removed_lines + + if change_type == "new": + if added_lines > 1: + final_code = "".join( + [ + line[1:] + for line in diff["diff"].split("\n") + if line.startswith("+") and not line.startswith("+++") + ] + ) + results.append( + { + "diff_url": diff_url, + "commit_sha": commit_sha, + "author_name": author_name, + "diff": final_code, + } + ) + else: + if total_changes > 1: + final_code = "".join( + [ + line[1:] + for line in diff["diff"].split("\n") + if (line.startswith("+") or line.startswith("-")) + and not line.startswith("+++") + and not line.startswith("---") + ] + ) + final_code_escaped = json.dumps(final_code)[1:-1] # Escape the final code + results.append( + { + "diff_url": diff_url, + "commit_sha": commit_sha, + "author_name": author_name, + "diff": final_code_escaped, + } + ) + except requests.RequestException as e: + print(f"Error fetching data from GitLab: {e}") + + return results diff --git a/api/core/tools/provider/builtin/gitlab/tools/gitlab_commits.yaml b/api/core/tools/provider/builtin/gitlab/tools/gitlab_commits.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2ff5fb570ecc42c5dd22420db406cef724626e8a --- /dev/null +++ b/api/core/tools/provider/builtin/gitlab/tools/gitlab_commits.yaml @@ -0,0 +1,88 @@ +identity: + name: gitlab_commits + author: Leo.Wang + label: + en_US: GitLab Commits + zh_Hans: GitLab 提交内容查询 +description: + human: + en_US: A tool for query GitLab commits, Input should be a exists username or project. + zh_Hans: 一个用于查询 GitLab 代码提交内容的工具,输入的内容应该是一个已存在的用户名或者项目名。 + llm: A tool for query GitLab commits, Input should be a exists username or project. +parameters: + - name: username + type: string + required: false + label: + en_US: username + zh_Hans: 员工用户名 + human_description: + en_US: username + zh_Hans: 员工用户名 + llm_description: User name for GitLab + form: llm + - name: repository + type: string + required: true + label: + en_US: repository + zh_Hans: 仓库路径 + human_description: + en_US: repository + zh_Hans: 仓库路径,以namespace/project_name的形式。 + llm_description: Repository path for GitLab, like namespace/project_name. + form: llm + - name: branch + type: string + required: false + label: + en_US: branch + zh_Hans: 分支名 + human_description: + en_US: branch + zh_Hans: 分支名 + llm_description: branch for GitLab + form: llm + - name: start_time + type: string + required: false + label: + en_US: start_time + zh_Hans: 开始时间 + human_description: + en_US: start_time + zh_Hans: 开始时间 + llm_description: Start time for GitLab + form: llm + - name: end_time + type: string + required: false + label: + en_US: end_time + zh_Hans: 结束时间 + human_description: + en_US: end_time + zh_Hans: 结束时间 + llm_description: End time for GitLab + form: llm + - name: change_type + type: select + required: false + options: + - value: all + label: + en_US: all + zh_Hans: 所有 + - value: new + label: + en_US: new + zh_Hans: 新增 + default: all + label: + en_US: change_type + zh_Hans: 变更类型 + human_description: + en_US: change_type + zh_Hans: 变更类型 + llm_description: Content change type for GitLab + form: llm diff --git a/api/core/tools/provider/builtin/gitlab/tools/gitlab_files.py b/api/core/tools/provider/builtin/gitlab/tools/gitlab_files.py new file mode 100644 index 0000000000000000000000000000000000000000..0ac9e2777dfe2c98e6ff6096da1eab7b86047e0b --- /dev/null +++ b/api/core/tools/provider/builtin/gitlab/tools/gitlab_files.py @@ -0,0 +1,103 @@ +import urllib.parse +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class GitlabFilesTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + repository = tool_parameters.get("repository", "") + project = tool_parameters.get("project", "") + branch = tool_parameters.get("branch", "") + path = tool_parameters.get("path", "") + file_path = tool_parameters.get("file_path", "") + + if not repository and not project: + return self.create_text_message("Either repository or project is required") + if not branch: + return self.create_text_message("Branch is required") + if not path and not file_path: + return self.create_text_message("Either path or file_path is required") + + access_token = self.runtime.credentials.get("access_tokens") + headers = {"PRIVATE-TOKEN": access_token} + site_url = self.runtime.credentials.get("site_url") + + if "access_tokens" not in self.runtime.credentials or not self.runtime.credentials.get("access_tokens"): + return self.create_text_message("Gitlab API Access Tokens is required.") + if "site_url" not in self.runtime.credentials or not self.runtime.credentials.get("site_url"): + site_url = "https://gitlab.com" + + if repository: + # URL encode the repository path + identifier = urllib.parse.quote(repository, safe="") + else: + identifier = self.get_project_id(site_url, access_token, project) + if not identifier: + raise Exception(f"Project '{project}' not found.)") + + # Get file content + if path: + results = self.fetch_files(site_url, headers, identifier, branch, path) + return [self.create_json_message(item) for item in results] + else: + result = self.fetch_file(site_url, headers, identifier, branch, file_path) + return [self.create_json_message(result)] + + @staticmethod + def fetch_file( + site_url: str, + headers: dict[str, str], + identifier: str, + branch: str, + path: str, + ) -> dict[str, Any]: + encoded_file_path = urllib.parse.quote(path, safe="") + file_url = f"{site_url}/api/v4/projects/{identifier}/repository/files/{encoded_file_path}/raw?ref={branch}" + + file_response = requests.get(file_url, headers=headers) + file_response.raise_for_status() + file_content = file_response.text + return {"path": path, "branch": branch, "content": file_content} + + def fetch_files( + self, site_url: str, headers: dict[str, str], identifier: str, branch: str, path: str + ) -> list[dict[str, Any]]: + results = [] + + try: + tree_url = f"{site_url}/api/v4/projects/{identifier}/repository/tree?path={path}&ref={branch}" + response = requests.get(tree_url, headers=headers) + response.raise_for_status() + items = response.json() + + for item in items: + item_path = item["path"] + if item["type"] == "tree": # It's a directory + results.extend(self.fetch_files(site_url, headers, identifier, branch, item_path)) + else: # It's a file + result = self.fetch_file(site_url, headers, identifier, branch, item_path) + results.append(result) + except requests.RequestException as e: + print(f"Error fetching data from GitLab: {e}") + + return results + + def get_project_id(self, site_url: str, access_token: str, project_name: str) -> Union[str, None]: + headers = {"PRIVATE-TOKEN": access_token} + try: + url = f"{site_url}/api/v4/projects?search={project_name}" + response = requests.get(url, headers=headers) + response.raise_for_status() + projects = response.json() + for project in projects: + if project["name"] == project_name: + return project["id"] + except requests.RequestException as e: + print(f"Error fetching project ID from GitLab: {e}") + return None diff --git a/api/core/tools/provider/builtin/gitlab/tools/gitlab_files.yaml b/api/core/tools/provider/builtin/gitlab/tools/gitlab_files.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3371f62fa8d98c5c11edc4c220f71775a3590e45 --- /dev/null +++ b/api/core/tools/provider/builtin/gitlab/tools/gitlab_files.yaml @@ -0,0 +1,65 @@ +identity: + name: gitlab_files + author: Leo.Wang + label: + en_US: GitLab Files + zh_Hans: GitLab 文件获取 +description: + human: + en_US: A tool for query GitLab files, Input should be branch and a exists file or directory path. + zh_Hans: 一个用于查询 GitLab 文件的工具,输入的内容应该是分支和一个已存在文件或者文件夹路径。 + llm: A tool for query GitLab files, Input should be a exists file or directory path. +parameters: + - name: repository + type: string + required: false + label: + en_US: repository + zh_Hans: 仓库路径 + human_description: + en_US: repository + zh_Hans: 仓库路径,以namespace/project_name的形式。 + llm_description: Repository path for GitLab, like namespace/project_name. + form: llm + - name: project + type: string + required: false + label: + en_US: project + zh_Hans: 项目 + human_description: + en_US: project + zh_Hans: 项目(和仓库路径二选一,都填写以仓库路径优先) + llm_description: Project for GitLab + form: llm + - name: branch + type: string + required: true + label: + en_US: branch + zh_Hans: 分支 + human_description: + en_US: branch + zh_Hans: 分支 + llm_description: Branch for GitLab + form: llm + - name: path + type: string + label: + en_US: path + zh_Hans: 文件夹 + human_description: + en_US: path + zh_Hans: 文件夹 + llm_description: Dir path for GitLab + form: llm + - name: file_path + type: string + label: + en_US: file_path + zh_Hans: 文件路径 + human_description: + en_US: file_path + zh_Hans: 文件路径(和文件夹二选一,都填写以文件夹优先) + llm_description: File path for GitLab + form: llm diff --git a/api/core/tools/provider/builtin/gitlab/tools/gitlab_mergerequests.py b/api/core/tools/provider/builtin/gitlab/tools/gitlab_mergerequests.py new file mode 100644 index 0000000000000000000000000000000000000000..ef99fa82e9d9d60277e06e49601a8b74c5b1536b --- /dev/null +++ b/api/core/tools/provider/builtin/gitlab/tools/gitlab_mergerequests.py @@ -0,0 +1,78 @@ +import urllib.parse +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class GitlabMergeRequestsTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + repository = tool_parameters.get("repository", "") + branch = tool_parameters.get("branch", "") + start_time = tool_parameters.get("start_time", "") + end_time = tool_parameters.get("end_time", "") + state = tool_parameters.get("state", "opened") # Default to "opened" + + if not repository: + return self.create_text_message("Repository is required") + + access_token = self.runtime.credentials.get("access_tokens") + site_url = self.runtime.credentials.get("site_url") + + if not access_token: + return self.create_text_message("Gitlab API Access Tokens is required.") + if not site_url: + site_url = "https://gitlab.com" + + # Get merge requests + result = self.get_merge_requests(site_url, access_token, repository, branch, start_time, end_time, state) + + return [self.create_json_message(item) for item in result] + + def get_merge_requests( + self, site_url: str, access_token: str, repository: str, branch: str, start_time: str, end_time: str, state: str + ) -> list[dict[str, Any]]: + domain = site_url + headers = {"PRIVATE-TOKEN": access_token} + results = [] + + try: + # URL encode the repository path + encoded_repository = urllib.parse.quote(repository, safe="") + merge_requests_url = f"{domain}/api/v4/projects/{encoded_repository}/merge_requests" + params = {"state": state} + + # Add time filters if provided + if start_time: + params["created_after"] = start_time + if end_time: + params["created_before"] = end_time + + response = requests.get(merge_requests_url, headers=headers, params=params) + response.raise_for_status() + merge_requests = response.json() + + for mr in merge_requests: + # Filter by target branch + if branch and mr["target_branch"] != branch: + continue + + results.append( + { + "id": mr["id"], + "title": mr["title"], + "author": mr["author"]["name"], + "web_url": mr["web_url"], + "target_branch": mr["target_branch"], + "created_at": mr["created_at"], + "state": mr["state"], + } + ) + except requests.RequestException as e: + print(f"Error fetching merge requests from GitLab: {e}") + + return results diff --git a/api/core/tools/provider/builtin/gitlab/tools/gitlab_mergerequests.yaml b/api/core/tools/provider/builtin/gitlab/tools/gitlab_mergerequests.yaml new file mode 100644 index 0000000000000000000000000000000000000000..81adb3db7d932d9ce06b48fab6b556ab8d2ae362 --- /dev/null +++ b/api/core/tools/provider/builtin/gitlab/tools/gitlab_mergerequests.yaml @@ -0,0 +1,77 @@ +identity: + name: gitlab_mergerequests + author: Leo.Wang + label: + en_US: GitLab Merge Requests + zh_Hans: GitLab 合并请求查询 +description: + human: + en_US: A tool for query GitLab merge requests, Input should be a exists repository or branch. + zh_Hans: 一个用于查询 GitLab 代码合并请求的工具,输入的内容应该是一个已存在的仓库名或者分支。 + llm: A tool for query GitLab merge requests, Input should be a exists repository or branch. +parameters: + - name: repository + type: string + required: false + label: + en_US: repository + zh_Hans: 仓库路径 + human_description: + en_US: repository + zh_Hans: 仓库路径,以namespace/project_name的形式。 + llm_description: Repository path for GitLab, like namespace/project_name. + form: llm + - name: branch + type: string + required: false + label: + en_US: branch + zh_Hans: 分支名 + human_description: + en_US: branch + zh_Hans: 分支名 + llm_description: branch for GitLab + form: llm + - name: start_time + type: string + required: false + label: + en_US: start_time + zh_Hans: 开始时间 + human_description: + en_US: start_time + zh_Hans: 开始时间 + llm_description: Start time for GitLab + form: llm + - name: end_time + type: string + required: false + label: + en_US: end_time + zh_Hans: 结束时间 + human_description: + en_US: end_time + zh_Hans: 结束时间 + llm_description: End time for GitLab + form: llm + - name: state + type: select + required: false + options: + - value: opened + label: + en_US: opened + zh_Hans: 打开 + - value: closed + label: + en_US: closed + zh_Hans: 关闭 + default: opened + label: + en_US: state + zh_Hans: 变更状态 + human_description: + en_US: state + zh_Hans: 变更状态 + llm_description: Merge request state type for GitLab + form: llm diff --git a/api/core/tools/provider/builtin/gitlab/tools/gitlab_projects.py b/api/core/tools/provider/builtin/gitlab/tools/gitlab_projects.py new file mode 100644 index 0000000000000000000000000000000000000000..ea0c028b4f3d079e2bc4ce530c57177ca8c79b9d --- /dev/null +++ b/api/core/tools/provider/builtin/gitlab/tools/gitlab_projects.py @@ -0,0 +1,81 @@ +import urllib.parse +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class GitlabProjectsTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + project_name = tool_parameters.get("project_name", "") + page = tool_parameters.get("page", 1) + page_size = tool_parameters.get("page_size", 20) + + access_token = self.runtime.credentials.get("access_tokens") + site_url = self.runtime.credentials.get("site_url") + + if not access_token: + return self.create_text_message("Gitlab API Access Tokens is required.") + if not site_url: + site_url = "https://gitlab.com" + + # Get project content + result = self.fetch_projects(site_url, access_token, project_name, page, page_size) + + return [self.create_json_message(item) for item in result] + + def fetch_projects( + self, + site_url: str, + access_token: str, + project_name: str, + page: str, + page_size: str, + ) -> list[dict[str, Any]]: + domain = site_url + headers = {"PRIVATE-TOKEN": access_token} + results = [] + + try: + if project_name: + # URL encode the project name for the search query + encoded_project_name = urllib.parse.quote(project_name, safe="") + projects_url = ( + f"{domain}/api/v4/projects?search={encoded_project_name}&page={page}&per_page={page_size}" + ) + else: + projects_url = f"{domain}/api/v4/projects?page={page}&per_page={page_size}" + + response = requests.get(projects_url, headers=headers) + response.raise_for_status() + projects = response.json() + + for project in projects: + # Filter projects by exact name match if necessary + if project_name and project["name"].lower() == project_name.lower(): + results.append( + { + "id": project["id"], + "name": project["name"], + "description": project.get("description", ""), + "web_url": project["web_url"], + } + ) + elif not project_name: + # If no specific project name is provided, add all projects + results.append( + { + "id": project["id"], + "name": project["name"], + "description": project.get("description", ""), + "web_url": project["web_url"], + } + ) + except requests.RequestException as e: + print(f"Error fetching data from GitLab: {e}") + + return results diff --git a/api/core/tools/provider/builtin/gitlab/tools/gitlab_projects.yaml b/api/core/tools/provider/builtin/gitlab/tools/gitlab_projects.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5fe098e1f7a647964e67e269ad82bbfb94471a3b --- /dev/null +++ b/api/core/tools/provider/builtin/gitlab/tools/gitlab_projects.yaml @@ -0,0 +1,45 @@ +identity: + name: gitlab_projects + author: Leo.Wang + label: + en_US: GitLab Projects + zh_Hans: GitLab 项目列表查询 +description: + human: + en_US: A tool for query GitLab projects, Input should be a project name. + zh_Hans: 一个用于查询 GitLab 项目列表的工具,输入的内容应该是一个项目名称。 + llm: A tool for query GitLab projects, Input should be a project name. +parameters: + - name: project_name + type: string + required: false + label: + en_US: project_name + zh_Hans: 项目名称 + human_description: + en_US: project_name + zh_Hans: 项目名称 + llm_description: Project name for GitLab + form: llm + - name: page + type: string + required: false + label: + en_US: page + zh_Hans: 页码 + human_description: + en_US: page + zh_Hans: 页码 + llm_description: Page index for GitLab + form: llm + - name: page_size + type: string + required: false + label: + en_US: page_size + zh_Hans: 每页数量 + human_description: + en_US: page_size + zh_Hans: 每页数量 + llm_description: Page size for GitLab + form: llm diff --git a/api/core/tools/provider/builtin/google/_assets/icon.svg b/api/core/tools/provider/builtin/google/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..bebbf52d3a23a45de00e86ce44dbf41252990117 --- /dev/null +++ b/api/core/tools/provider/builtin/google/_assets/icon.svg @@ -0,0 +1,6 @@ + + + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/google/google.py b/api/core/tools/provider/builtin/google/google.py new file mode 100644 index 0000000000000000000000000000000000000000..6b5395f9d3e5b82162c42aca21b94f46d7f42c96 --- /dev/null +++ b/api/core/tools/provider/builtin/google/google.py @@ -0,0 +1,20 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.google.tools.google_search import GoogleSearchTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class GoogleProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + GoogleSearchTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={"query": "test", "result_type": "link"}, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/google/google.yaml b/api/core/tools/provider/builtin/google/google.yaml new file mode 100644 index 0000000000000000000000000000000000000000..afb4d5b2145ba6c706ae7c1c38fcb67fff4a586b --- /dev/null +++ b/api/core/tools/provider/builtin/google/google.yaml @@ -0,0 +1,31 @@ +identity: + author: Dify + name: google + label: + en_US: Google + zh_Hans: Google + pt_BR: Google + description: + en_US: Google + zh_Hans: GoogleSearch + pt_BR: Google + icon: icon.svg + tags: + - search +credentials_for_provider: + serpapi_api_key: + type: secret-input + required: true + label: + en_US: SerpApi API key + zh_Hans: SerpApi API key + pt_BR: SerpApi API key + placeholder: + en_US: Please input your SerpApi API key + zh_Hans: 请输入你的 SerpApi API key + pt_BR: Please input your SerpApi API key + help: + en_US: Get your SerpApi API key from SerpApi + zh_Hans: 从 SerpApi 获取您的 SerpApi API key + pt_BR: Get your SerpApi API key from SerpApi + url: https://serpapi.com/manage-api-key diff --git a/api/core/tools/provider/builtin/google/tools/google_search.py b/api/core/tools/provider/builtin/google/tools/google_search.py new file mode 100644 index 0000000000000000000000000000000000000000..a9f65925d86f9425f12caca8deba52daca075758 --- /dev/null +++ b/api/core/tools/provider/builtin/google/tools/google_search.py @@ -0,0 +1,40 @@ +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +SERP_API_URL = "https://serpapi.com/search" + + +class GoogleSearchTool(BuiltinTool): + def _parse_response(self, response: dict) -> dict: + result = {} + if "knowledge_graph" in response: + result["title"] = response["knowledge_graph"].get("title", "") + result["description"] = response["knowledge_graph"].get("description", "") + if "organic_results" in response: + result["organic_results"] = [ + {"title": item.get("title", ""), "link": item.get("link", ""), "snippet": item.get("snippet", "")} + for item in response["organic_results"] + ] + return result + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + params = { + "api_key": self.runtime.credentials["serpapi_api_key"], + "q": tool_parameters["query"], + "engine": "google", + "google_domain": "google.com", + "gl": "us", + "hl": "en", + } + response = requests.get(url=SERP_API_URL, params=params) + response.raise_for_status() + valuable_res = self._parse_response(response.json()) + return self.create_json_message(valuable_res) diff --git a/api/core/tools/provider/builtin/google/tools/google_search.yaml b/api/core/tools/provider/builtin/google/tools/google_search.yaml new file mode 100644 index 0000000000000000000000000000000000000000..72db3839eb022ae557f4fd3cc1c8430b520403b1 --- /dev/null +++ b/api/core/tools/provider/builtin/google/tools/google_search.yaml @@ -0,0 +1,27 @@ +identity: + name: google_search + author: Dify + label: + en_US: GoogleSearch + zh_Hans: 谷歌搜索 + pt_BR: GoogleSearch +description: + human: + en_US: A tool for performing a Google SERP search and extracting snippets and webpages.Input should be a search query. + zh_Hans: 一个用于执行 Google SERP 搜索并提取片段和网页的工具。输入应该是一个搜索查询。 + pt_BR: A tool for performing a Google SERP search and extracting snippets and webpages.Input should be a search query. + llm: A tool for performing a Google SERP search and extracting snippets and webpages.Input should be a search query. +parameters: + - name: query + type: string + required: true + label: + en_US: Query string + zh_Hans: 查询语句 + pt_BR: Query string + human_description: + en_US: used for searching + zh_Hans: 用于搜索网页内容 + pt_BR: used for searching + llm_description: key words for searching + form: llm diff --git a/api/core/tools/provider/builtin/google_translate/_assets/icon.svg b/api/core/tools/provider/builtin/google_translate/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..de69a9c5e583162ca0abcf6faa651742b7745922 --- /dev/null +++ b/api/core/tools/provider/builtin/google_translate/_assets/icon.svg @@ -0,0 +1,18 @@ + + + + + + + + + + + + + + + + + + diff --git a/api/core/tools/provider/builtin/google_translate/google_translate.py b/api/core/tools/provider/builtin/google_translate/google_translate.py new file mode 100644 index 0000000000000000000000000000000000000000..ea53aa4eeb906ff6fb2b7c3e3fc47b39cdfd7976 --- /dev/null +++ b/api/core/tools/provider/builtin/google_translate/google_translate.py @@ -0,0 +1,13 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.google_translate.tools.translate import GoogleTranslate +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class JsonExtractProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + GoogleTranslate().invoke(user_id="", tool_parameters={"content": "这是一段测试文本", "dest": "en"}) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/google_translate/google_translate.yaml b/api/core/tools/provider/builtin/google_translate/google_translate.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8bc821a3d5e9faf853690e09cdfd72843f2e2d9c --- /dev/null +++ b/api/core/tools/provider/builtin/google_translate/google_translate.yaml @@ -0,0 +1,12 @@ +identity: + author: Ron Liu + name: google_translate + label: + en_US: Google Translate + zh_Hans: 谷歌翻译 + description: + en_US: Translate text using Google + zh_Hans: 使用 Google 进行翻译 + icon: icon.svg + tags: + - utilities diff --git a/api/core/tools/provider/builtin/google_translate/tools/translate.py b/api/core/tools/provider/builtin/google_translate/tools/translate.py new file mode 100644 index 0000000000000000000000000000000000000000..ea3f2077d5d4855a15cb44694fe75e114d6c098e --- /dev/null +++ b/api/core/tools/provider/builtin/google_translate/tools/translate.py @@ -0,0 +1,47 @@ +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class GoogleTranslate(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + content = tool_parameters.get("content", "") + if not content: + return self.create_text_message("Invalid parameter content") + + dest = tool_parameters.get("dest", "") + if not dest: + return self.create_text_message("Invalid parameter destination language") + + try: + result = self._translate(content, dest) + return self.create_text_message(str(result)) + except Exception: + return self.create_text_message("Translation service error, please check the network") + + def _translate(self, content: str, dest: str) -> str: + try: + url = "https://translate.googleapis.com/translate_a/single" + params = {"client": "gtx", "sl": "auto", "tl": dest, "dt": "t", "q": content} + + headers = { + "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko)" + " Chrome/91.0.4472.124 Safari/537.36" + } + + response_json = requests.get(url, params=params, headers=headers).json() + result = response_json[0] + translated_text = "".join([item[0] for item in result if item[0]]) + return str(translated_text) + except Exception as e: + return str(e) diff --git a/api/core/tools/provider/builtin/google_translate/tools/translate.yaml b/api/core/tools/provider/builtin/google_translate/tools/translate.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a4189cd7439ad761768bfc51f2225cf85ca8267e --- /dev/null +++ b/api/core/tools/provider/builtin/google_translate/tools/translate.yaml @@ -0,0 +1,215 @@ +identity: + name: translate + author: Ron Liu + label: + en_US: Translate + zh_Hans: 翻译 +description: + human: + en_US: A tool for Google Translate + zh_Hans: Google 翻译 + llm: A tool for Google Translate +parameters: + - name: content + type: string + required: true + label: + en_US: Text content + zh_Hans: 文本内容 + human_description: + en_US: Text content + zh_Hans: 需要翻译的文本内容 + llm_description: Text content + form: llm + - name: dest + type: select + required: true + label: + en_US: destination language + zh_Hans: 目标语言 + human_description: + en_US: The destination language you want to translate. + zh_Hans: 你想翻译的目标语言 + default: en + form: form + options: + - value: ar + label: + en_US: Arabic + zh_Hans: 阿拉伯语 + - value: bg + label: + en_US: Bulgarian + zh_Hans: 保加利亚语 + - value: ca + label: + en_US: Catalan + zh_Hans: 加泰罗尼亚语 + - value: zh-cn + label: + en_US: Chinese (Simplified) + zh_Hans: 中文(简体) + - value: zh-tw + label: + en_US: Chinese (Traditional) + zh_Hans: 中文(繁体) + - value: cs + label: + en_US: Czech + zh_Hans: 捷克语 + - value: da + label: + en_US: Danish + zh_Hans: 丹麦语 + - value: nl + label: + en_US: Dutch + zh_Hans: 荷兰语 + - value: en + label: + en_US: English + zh_Hans: 英语 + - value: et + label: + en_US: Estonian + zh_Hans: 爱沙尼亚语 + - value: fi + label: + en_US: Finnish + zh_Hans: 芬兰语 + - value: fr + label: + en_US: French + zh_Hans: 法语 + - value: de + label: + en_US: German + zh_Hans: 德语 + - value: el + label: + en_US: Greek + zh_Hans: 希腊语 + - value: iw + label: + en_US: Hebrew + zh_Hans: 希伯来语 + - value: hi + label: + en_US: Hindi + zh_Hans: 印地语 + - value: hu + label: + en_US: Hungarian + zh_Hans: 匈牙利语 + - value: id + label: + en_US: Indonesian + zh_Hans: 印尼语 + - value: it + label: + en_US: Italian + zh_Hans: 意大利语 + - value: ja + label: + en_US: Japanese + zh_Hans: 日语 + - value: kn + label: + en_US: Kannada + zh_Hans: 卡纳达语 + - value: ko + label: + en_US: Korean + zh_Hans: 韩语 + - value: lv + label: + en_US: Latvian + zh_Hans: 拉脱维亚语 + - value: lt + label: + en_US: Lithuanian + zh_Hans: 立陶宛语 + - value: my + label: + en_US: Malay + zh_Hans: 马来语 + - value: ml + label: + en_US: Malayalam + zh_Hans: 马拉雅拉姆语 + - value: mr + label: + en_US: Marathi + zh_Hans: 马拉地语 + - value: "no" + label: + en_US: Norwegian + zh_Hans: 挪威语 + - value: pl + label: + en_US: Polish + zh_Hans: 波兰语 + - value: pt-br + label: + en_US: Portuguese (Brazil) + zh_Hans: 葡萄牙语(巴西) + - value: pt-pt + label: + en_US: Portuguese (Portugal) + zh_Hans: 葡萄牙语(葡萄牙) + - value: pa + label: + en_US: Punjabi + zh_Hans: 旁遮普语 + - value: ro + label: + en_US: Romanian + zh_Hans: 罗马尼亚语 + - value: ru + label: + en_US: Russian + zh_Hans: 俄语 + - value: sr + label: + en_US: Serbian + zh_Hans: 塞尔维亚语 + - value: sk + label: + en_US: Slovak + zh_Hans: 斯洛伐克语 + - value: sl + label: + en_US: Slovenian + zh_Hans: 斯洛文尼亚语 + - value: es + label: + en_US: Spanish + zh_Hans: 西班牙语 + - value: sv + label: + en_US: Swedish + zh_Hans: 瑞典语 + - value: ta + label: + en_US: Tamil + zh_Hans: 泰米尔语 + - value: te + label: + en_US: Telugu + zh_Hans: 泰卢固语 + - value: th + label: + en_US: Thai + zh_Hans: 泰语 + - value: tr + label: + en_US: Turkish + zh_Hans: 土耳其语 + - value: uk + label: + en_US: Ukrainian + zh_Hans: 乌克兰语 + - value: vi + label: + en_US: Vietnamese + zh_Hans: 越南语 diff --git a/api/core/tools/provider/builtin/hap/_assets/icon.svg b/api/core/tools/provider/builtin/hap/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..0fa6f0886fdfdb35d287d46f722c5288b1cd97ca --- /dev/null +++ b/api/core/tools/provider/builtin/hap/_assets/icon.svg @@ -0,0 +1,16 @@ + + + + + + + + + + + + + + + + diff --git a/api/core/tools/provider/builtin/hap/hap.py b/api/core/tools/provider/builtin/hap/hap.py new file mode 100644 index 0000000000000000000000000000000000000000..cbdf95046595687501a64f888a5cfe2cf7fa528f --- /dev/null +++ b/api/core/tools/provider/builtin/hap/hap.py @@ -0,0 +1,8 @@ +from typing import Any + +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class HapProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + pass diff --git a/api/core/tools/provider/builtin/hap/hap.yaml b/api/core/tools/provider/builtin/hap/hap.yaml new file mode 100644 index 0000000000000000000000000000000000000000..25b473cf9dd21199d4fbb82337a422c63acf3eaf --- /dev/null +++ b/api/core/tools/provider/builtin/hap/hap.yaml @@ -0,0 +1,15 @@ +identity: + author: Mingdao + name: hap + label: + en_US: HAP + zh_Hans: HAP + pt_BR: HAP + description: + en_US: "Hyper application platform that is particularly friendly to AI" + zh_Hans: "对 AI 特别友好的超级应用平台" + pt_BR: "Plataforma de aplicação hiper que é particularmente amigável à IA" + icon: icon.svg + tags: + - productivity +credentials_for_provider: diff --git a/api/core/tools/provider/builtin/hap/tools/add_worksheet_record.py b/api/core/tools/provider/builtin/hap/tools/add_worksheet_record.py new file mode 100644 index 0000000000000000000000000000000000000000..597adc91db9768256d99cb89ce645caaf391aa5c --- /dev/null +++ b/api/core/tools/provider/builtin/hap/tools/add_worksheet_record.py @@ -0,0 +1,52 @@ +import json +from typing import Any, Union + +import httpx + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class AddWorksheetRecordTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + appkey = tool_parameters.get("appkey", "") + if not appkey: + return self.create_text_message("Invalid parameter App Key") + sign = tool_parameters.get("sign", "") + if not sign: + return self.create_text_message("Invalid parameter Sign") + worksheet_id = tool_parameters.get("worksheet_id", "") + if not worksheet_id: + return self.create_text_message("Invalid parameter Worksheet ID") + record_data = tool_parameters.get("record_data", "") + if not record_data: + return self.create_text_message("Invalid parameter Record Row Data") + + host = tool_parameters.get("host", "") + if not host: + host = "https://api.mingdao.com" + elif not host.startswith(("http://", "https://")): + return self.create_text_message("Invalid parameter Host Address") + else: + host = f"{host.removesuffix('/')}/api" + + url = f"{host}/v2/open/worksheet/addRow" + headers = {"Content-Type": "application/json"} + payload = {"appKey": appkey, "sign": sign, "worksheetId": worksheet_id} + + try: + payload["controls"] = json.loads(record_data) + res = httpx.post(url, headers=headers, json=payload, timeout=60) + res.raise_for_status() + res_json = res.json() + if res_json.get("error_code") != 1: + return self.create_text_message(f"Failed to add the new record. {res_json['error_msg']}") + return self.create_text_message(f"New record added successfully. The record ID is {res_json['data']}.") + except httpx.RequestError as e: + return self.create_text_message(f"Failed to add the new record, request error: {e}") + except json.JSONDecodeError as e: + return self.create_text_message(f"Failed to parse JSON response: {e}") + except Exception as e: + return self.create_text_message(f"Failed to add the new record, unexpected error: {e}") diff --git a/api/core/tools/provider/builtin/hap/tools/add_worksheet_record.yaml b/api/core/tools/provider/builtin/hap/tools/add_worksheet_record.yaml new file mode 100644 index 0000000000000000000000000000000000000000..add7742cd74db1cb597cf94cb4cebc98bf6ba50b --- /dev/null +++ b/api/core/tools/provider/builtin/hap/tools/add_worksheet_record.yaml @@ -0,0 +1,78 @@ +identity: + name: add_worksheet_record + author: Ryan Tian + label: + en_US: Add Worksheet Record + zh_Hans: 新增一条工作表记录 +description: + human: + en_US: Adds a new record to the specified worksheet + zh_Hans: 向指定的工作表新增一条记录数据 + llm: A tool to append a new data entry into a specified worksheet. +parameters: + - name: appkey + type: secret-input + required: true + label: + en_US: App Key + zh_Hans: App Key + human_description: + en_US: The AppKey parameter for the HAP application, typically found in the application's API documentation. + zh_Hans: HAP 应用的 AppKey 参数,可以从应用 API 文档中查找到 + llm_description: the AppKey parameter for the HAP application + form: form + + - name: sign + type: secret-input + required: true + label: + en_US: Sign + zh_Hans: Sign + human_description: + en_US: The Sign parameter for the HAP application + zh_Hans: HAP 应用的 Sign 参数 + llm_description: the Sign parameter for the HAP application + form: form + + - name: worksheet_id + type: string + required: true + label: + en_US: Worksheet ID + zh_Hans: 工作表 ID + human_description: + en_US: The ID of the specified worksheet + zh_Hans: 要获取字段信息的工作表 ID + llm_description: The ID of the specified worksheet which to get the fields information. + form: llm + + - name: record_data + type: string + required: true + label: + en_US: Record Row Data + zh_Hans: 记录数据 + human_description: + en_US: The fields with data of the specified record + zh_Hans: 要新增的记录数据,JSON 对象数组格式。数组元素属性:controlId-字段ID,value-字段值 + llm_description: | + The fields with data of the specified record which to be created. It is in the format of an array of JSON objects, and the structure is defined as follows: + ``` + type RowData = { + controlId: string; // Field ID to be updated + value: string; // Field value to be updated + }[]; + ``` + form: llm + + - name: host + type: string + required: false + label: + en_US: Host Address + zh_Hans: 服务器地址 + human_description: + en_US: The address for the privately deployed HAP server. + zh_Hans: 私有部署 HAP 服务器地址,公有云无需填写 + llm_description: the address for the privately deployed HAP server. + form: form diff --git a/api/core/tools/provider/builtin/hap/tools/delete_worksheet_record.py b/api/core/tools/provider/builtin/hap/tools/delete_worksheet_record.py new file mode 100644 index 0000000000000000000000000000000000000000..5d42af4c490598585fd3f3bf08bf56b0712a1d3f --- /dev/null +++ b/api/core/tools/provider/builtin/hap/tools/delete_worksheet_record.py @@ -0,0 +1,48 @@ +from typing import Any, Union + +import httpx + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class DeleteWorksheetRecordTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + appkey = tool_parameters.get("appkey", "") + if not appkey: + return self.create_text_message("Invalid parameter App Key") + sign = tool_parameters.get("sign", "") + if not sign: + return self.create_text_message("Invalid parameter Sign") + worksheet_id = tool_parameters.get("worksheet_id", "") + if not worksheet_id: + return self.create_text_message("Invalid parameter Worksheet ID") + row_id = tool_parameters.get("row_id", "") + if not row_id: + return self.create_text_message("Invalid parameter Record Row ID") + + host = tool_parameters.get("host", "") + if not host: + host = "https://api.mingdao.com" + elif not host.startswith(("http://", "https://")): + return self.create_text_message("Invalid parameter Host Address") + else: + host = f"{host.removesuffix('/')}/api" + + url = f"{host}/v2/open/worksheet/deleteRow" + headers = {"Content-Type": "application/json"} + payload = {"appKey": appkey, "sign": sign, "worksheetId": worksheet_id, "rowId": row_id} + + try: + res = httpx.post(url, headers=headers, json=payload, timeout=30) + res.raise_for_status() + res_json = res.json() + if res_json.get("error_code") != 1: + return self.create_text_message(f"Failed to delete the record. {res_json['error_msg']}") + return self.create_text_message("Successfully deleted the record.") + except httpx.RequestError as e: + return self.create_text_message(f"Failed to delete the record, request error: {e}") + except Exception as e: + return self.create_text_message(f"Failed to delete the record, unexpected error: {e}") diff --git a/api/core/tools/provider/builtin/hap/tools/delete_worksheet_record.yaml b/api/core/tools/provider/builtin/hap/tools/delete_worksheet_record.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7c0c2a6439003f6795a08438aef13569dcc3bd7a --- /dev/null +++ b/api/core/tools/provider/builtin/hap/tools/delete_worksheet_record.yaml @@ -0,0 +1,71 @@ +identity: + name: delete_worksheet_record + author: Ryan Tian + label: + en_US: Delete Worksheet Record + zh_Hans: 删除指定的一条工作表记录 +description: + human: + en_US: Deletes a single record from a worksheet based on the specified record row ID + zh_Hans: 根据指定的记录ID删除一条工作表记录数据 + llm: A tool to remove a particular record from a worksheet by specifying its unique record identifier. +parameters: + - name: appkey + type: secret-input + required: true + label: + en_US: App Key + zh_Hans: App Key + human_description: + en_US: The AppKey parameter for the HAP application, typically found in the application's API documentation. + zh_Hans: HAP 应用的 AppKey 参数,可以从应用 API 文档中查找到 + llm_description: the AppKey parameter for the HAP application + form: form + + - name: sign + type: secret-input + required: true + label: + en_US: Sign + zh_Hans: Sign + human_description: + en_US: The Sign parameter for the HAP application + zh_Hans: HAP 应用的 Sign 参数 + llm_description: the Sign parameter for the HAP application + form: form + + - name: worksheet_id + type: string + required: true + label: + en_US: Worksheet ID + zh_Hans: 工作表 ID + human_description: + en_US: The ID of the specified worksheet + zh_Hans: 要获取字段信息的工作表 ID + llm_description: The ID of the specified worksheet which to get the fields information. + form: llm + + - name: row_id + type: string + required: true + label: + en_US: Record Row ID + zh_Hans: 记录 ID + human_description: + en_US: The row ID of the specified record + zh_Hans: 要删除的记录 ID + llm_description: The row ID of the specified record which to be deleted. + form: llm + + - name: host + type: string + required: false + label: + en_US: Host Address + zh_Hans: 服务器地址 + human_description: + en_US: The address for the privately deployed HAP server. + zh_Hans: 私有部署 HAP 服务器地址,公有云无需填写 + llm_description: the address for the privately deployed HAP server. + form: form diff --git a/api/core/tools/provider/builtin/hap/tools/get_worksheet_fields.py b/api/core/tools/provider/builtin/hap/tools/get_worksheet_fields.py new file mode 100644 index 0000000000000000000000000000000000000000..6887b8b4e99df68e4594eff57f0106cff0701d3e --- /dev/null +++ b/api/core/tools/provider/builtin/hap/tools/get_worksheet_fields.py @@ -0,0 +1,152 @@ +import json +from typing import Any, Union + +import httpx + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class GetWorksheetFieldsTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + appkey = tool_parameters.get("appkey", "") + if not appkey: + return self.create_text_message("Invalid parameter App Key") + sign = tool_parameters.get("sign", "") + if not sign: + return self.create_text_message("Invalid parameter Sign") + worksheet_id = tool_parameters.get("worksheet_id", "") + if not worksheet_id: + return self.create_text_message("Invalid parameter Worksheet ID") + + host = tool_parameters.get("host", "") + if not host: + host = "https://api.mingdao.com" + elif not host.startswith(("http://", "https://")): + return self.create_text_message("Invalid parameter Host Address") + else: + host = f"{host.removesuffix('/')}/api" + + url = f"{host}/v2/open/worksheet/getWorksheetInfo" + headers = {"Content-Type": "application/json"} + payload = {"appKey": appkey, "sign": sign, "worksheetId": worksheet_id} + + try: + res = httpx.post(url, headers=headers, json=payload, timeout=60) + res.raise_for_status() + res_json = res.json() + if res_json.get("error_code") != 1: + return self.create_text_message(f"Failed to get the worksheet information. {res_json['error_msg']}") + + fields_json, fields_table = self.get_controls(res_json["data"]["controls"]) + result_type = tool_parameters.get("result_type", "table") + return self.create_text_message( + text=json.dumps(fields_json, ensure_ascii=False) if result_type == "json" else fields_table + ) + except httpx.RequestError as e: + return self.create_text_message(f"Failed to get the worksheet information, request error: {e}") + except json.JSONDecodeError as e: + return self.create_text_message(f"Failed to parse JSON response: {e}") + except Exception as e: + return self.create_text_message(f"Failed to get the worksheet information, unexpected error: {e}") + + def get_field_type_by_id(self, field_type_id: int) -> str: + field_type_map = { + 2: "Text", + 3: "Text-Phone", + 4: "Text-Phone", + 5: "Text-Email", + 6: "Number", + 7: "Text", + 8: "Number", + 9: "Option-Single Choice", + 10: "Option-Multiple Choices", + 11: "Option-Single Choice", + 15: "Date", + 16: "Date", + 24: "Option-Region", + 25: "Text", + 26: "Option-Member", + 27: "Option-Department", + 28: "Number", + 29: "Option-Linked Record", + 30: "Unknown Type", + 31: "Number", + 32: "Text", + 33: "Text", + 35: "Option-Linked Record", + 36: "Number-Yes1/No0", + 37: "Number", + 38: "Date", + 40: "Location", + 41: "Text", + 46: "Time", + 48: "Option-Organizational Role", + 50: "Text", + 51: "Query Record", + } + return field_type_map.get(field_type_id, "") + + def get_controls(self, controls: list) -> dict: + fields = [] + fields_list = ["|fieldId|fieldName|fieldType|fieldTypeId|description|options|", "|" + "---|" * 6] + for control in controls: + if control["type"] in self._get_ignore_types(): + continue + field_type_id = control["type"] + field_type = self.get_field_type_by_id(control["type"]) + if field_type_id == 30: + source_type = control["sourceControl"]["type"] + if source_type in self._get_ignore_types(): + continue + else: + field_type_id = source_type + field_type = self.get_field_type_by_id(source_type) + field = { + "id": control["controlId"], + "name": control["controlName"], + "type": field_type, + "typeId": field_type_id, + "description": control["remark"].replace("\n", " ").replace("\t", " "), + "options": self._extract_options(control), + } + fields.append(field) + fields_list.append( + f"|{field['id']}|{field['name']}|{field['type']}|{field['typeId']}|{field['description']}" + f"|{field['options'] or ''}|" + ) + + fields.append( + { + "id": "ctime", + "name": "Created Time", + "type": self.get_field_type_by_id(16), + "typeId": 16, + "description": "", + "options": [], + } + ) + fields_list.append("|ctime|Created Time|Date|16|||") + return fields, "\n".join(fields_list) + + def _extract_options(self, control: dict) -> list: + options = [] + if control["type"] in {9, 10, 11}: + options.extend([{"key": opt["key"], "value": opt["value"]} for opt in control.get("options", [])]) + elif control["type"] in {28, 36}: + itemnames = control["advancedSetting"].get("itemnames") + if itemnames and itemnames.startswith("[{"): + try: + options = json.loads(itemnames) + except json.JSONDecodeError: + pass + elif control["type"] == 30: + source_type = control["sourceControl"]["type"] + if source_type not in self._get_ignore_types(): + options.extend([{"key": opt["key"], "value": opt["value"]} for opt in control.get("options", [])]) + return options + + def _get_ignore_types(self): + return {14, 21, 22, 34, 42, 43, 45, 47, 49, 10010} diff --git a/api/core/tools/provider/builtin/hap/tools/get_worksheet_fields.yaml b/api/core/tools/provider/builtin/hap/tools/get_worksheet_fields.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f0d4973e8549f7bdbdeb2a23a95a6cdf0d4923b6 --- /dev/null +++ b/api/core/tools/provider/builtin/hap/tools/get_worksheet_fields.yaml @@ -0,0 +1,80 @@ +identity: + name: get_worksheet_fields + author: Ryan Tian + label: + en_US: Get Worksheet Fields + zh_Hans: 获取工作表字段结构 +description: + human: + en_US: Get fields information of the worksheet + zh_Hans: 获取指定工作表的所有字段结构信息 + llm: A tool to get fields information of the specific worksheet. +parameters: + - name: appkey + type: secret-input + required: true + label: + en_US: App Key + zh_Hans: App Key + human_description: + en_US: The AppKey parameter for the HAP application, typically found in the application's API documentation. + zh_Hans: HAP 应用的 AppKey 参数,可以从应用 API 文档中查找到 + llm_description: the AppKey parameter for the HAP application + form: form + + - name: sign + type: secret-input + required: true + label: + en_US: Sign + zh_Hans: Sign + human_description: + en_US: The Sign parameter for the HAP application + zh_Hans: HAP 应用的 Sign 参数 + llm_description: the Sign parameter for the HAP application + form: form + + - name: worksheet_id + type: string + required: true + label: + en_US: Worksheet ID + zh_Hans: 工作表 ID + human_description: + en_US: The ID of the specified worksheet + zh_Hans: 要获取字段信息的工作表 ID + llm_description: The ID of the specified worksheet which to get the fields information. + form: llm + + - name: host + type: string + required: false + label: + en_US: Host Address + zh_Hans: 服务器地址 + human_description: + en_US: The address for the privately deployed HAP server. + zh_Hans: 私有部署 HAP 服务器地址,公有云无需填写 + llm_description: the address for the privately deployed HAP server. + form: form + + - name: result_type + type: select + required: true + options: + - value: table + label: + en_US: table text + zh_Hans: 表格文本 + - value: json + label: + en_US: json text + zh_Hans: JSON文本 + default: table + label: + en_US: Result type + zh_Hans: 结果类型 + human_description: + en_US: used for selecting the result type, table styled text or json text + zh_Hans: 用于选择结果类型,使用表格格式文本还是JSON格式文本 + form: form diff --git a/api/core/tools/provider/builtin/hap/tools/get_worksheet_pivot_data.py b/api/core/tools/provider/builtin/hap/tools/get_worksheet_pivot_data.py new file mode 100644 index 0000000000000000000000000000000000000000..26d7116869b6d9ca26c2541b60fee14621a3fa6a --- /dev/null +++ b/api/core/tools/provider/builtin/hap/tools/get_worksheet_pivot_data.py @@ -0,0 +1,137 @@ +import json +from typing import Any, Union + +import httpx + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class GetWorksheetPivotDataTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + appkey = tool_parameters.get("appkey", "") + if not appkey: + return self.create_text_message("Invalid parameter App Key") + sign = tool_parameters.get("sign", "") + if not sign: + return self.create_text_message("Invalid parameter Sign") + worksheet_id = tool_parameters.get("worksheet_id", "") + if not worksheet_id: + return self.create_text_message("Invalid parameter Worksheet ID") + x_column_fields = tool_parameters.get("x_column_fields", "") + if not x_column_fields or not x_column_fields.startswith("["): + return self.create_text_message("Invalid parameter Column Fields") + y_row_fields = tool_parameters.get("y_row_fields", "") + if y_row_fields and not y_row_fields.strip().startswith("["): + return self.create_text_message("Invalid parameter Row Fields") + elif not y_row_fields: + y_row_fields = "[]" + value_fields = tool_parameters.get("value_fields", "") + if not value_fields or not value_fields.strip().startswith("["): + return self.create_text_message("Invalid parameter Value Fields") + + host = tool_parameters.get("host", "") + if not host: + host = "https://api.mingdao.com" + elif not host.startswith(("http://", "https://")): + return self.create_text_message("Invalid parameter Host Address") + else: + host = f"{host.removesuffix('/')}/api" + + url = f"{host}/report/getPivotData" + headers = {"Content-Type": "application/json"} + payload = {"appKey": appkey, "sign": sign, "worksheetId": worksheet_id, "options": {"showTotal": True}} + + try: + x_column_fields = json.loads(x_column_fields) + payload["columns"] = x_column_fields + y_row_fields = json.loads(y_row_fields) + if y_row_fields: + payload["rows"] = y_row_fields + value_fields = json.loads(value_fields) + payload["values"] = value_fields + sort_fields = tool_parameters.get("sort_fields", "") + if not sort_fields: + sort_fields = "[]" + sort_fields = json.loads(sort_fields) + if sort_fields: + payload["options"]["sort"] = sort_fields + res = httpx.post(url, headers=headers, json=payload, timeout=60) + res.raise_for_status() + res_json = res.json() + if res_json.get("status") != 1: + return self.create_text_message(f"Failed to get the worksheet pivot data. {res_json['msg']}") + + pivot_json = self.generate_pivot_json(res_json["data"]) + pivot_table = self.generate_pivot_table(res_json["data"]) + result_type = tool_parameters.get("result_type", "") + text = pivot_table if result_type == "table" else json.dumps(pivot_json, ensure_ascii=False) + return self.create_text_message(text) + except httpx.RequestError as e: + return self.create_text_message(f"Failed to get the worksheet pivot data, request error: {e}") + except json.JSONDecodeError as e: + return self.create_text_message(f"Failed to parse JSON response: {e}") + except Exception as e: + return self.create_text_message(f"Failed to get the worksheet pivot data, unexpected error: {e}") + + def generate_pivot_table(self, data: dict[str, Any]) -> str: + columns = data["metadata"]["columns"] + rows = data["metadata"]["rows"] + values = data["metadata"]["values"] + + rows_data = data["data"] + + header = ( + ([row["displayName"] for row in rows] if rows else []) + + [column["displayName"] for column in columns] + + [value["displayName"] for value in values] + ) + line = (["---"] * len(rows) if rows else []) + ["---"] * len(columns) + ["--:"] * len(values) + + table = [header, line] + for row in rows_data: + row_data = [self.replace_pipe(row["rows"][r["controlId"]]) for r in rows] if rows else [] + row_data.extend([self.replace_pipe(row["columns"][column["controlId"]]) for column in columns]) + row_data.extend([self.replace_pipe(str(row["values"][value["controlId"]])) for value in values]) + table.append(row_data) + + return "\n".join([("|" + "|".join(row) + "|") for row in table]) + + def replace_pipe(self, text: str) -> str: + return text.replace("|", "▏").replace("\n", " ") + + def generate_pivot_json(self, data: dict[str, Any]) -> dict: + fields = { + "x-axis": [ + {"fieldId": column["controlId"], "fieldName": column["displayName"]} + for column in data["metadata"]["columns"] + ], + "y-axis": [ + {"fieldId": row["controlId"], "fieldName": row["displayName"]} for row in data["metadata"]["rows"] + ] + if data["metadata"]["rows"] + else [], + "values": [ + {"fieldId": value["controlId"], "fieldName": value["displayName"]} + for value in data["metadata"]["values"] + ], + } + # fields = ([ + # {"fieldId": row["controlId"], "fieldName": row["displayName"]} + # for row in data["metadata"]["rows"] + # ] if data["metadata"]["rows"] else []) + [ + # {"fieldId": column["controlId"], "fieldName": column["displayName"]} + # for column in data["metadata"]["columns"] + # ] + [ + # {"fieldId": value["controlId"], "fieldName": value["displayName"]} + # for value in data["metadata"]["values"] + # ] + rows = [] + for row in data["data"]: + row_data = row["rows"] or {} + row_data.update(row["columns"]) + row_data.update(row["values"]) + rows.append(row_data) + return {"fields": fields, "rows": rows, "summary": data["metadata"]["totalRow"]} diff --git a/api/core/tools/provider/builtin/hap/tools/get_worksheet_pivot_data.yaml b/api/core/tools/provider/builtin/hap/tools/get_worksheet_pivot_data.yaml new file mode 100644 index 0000000000000000000000000000000000000000..cf8c57b26208a915786b5e8b6ed4e4f4b8321aab --- /dev/null +++ b/api/core/tools/provider/builtin/hap/tools/get_worksheet_pivot_data.yaml @@ -0,0 +1,248 @@ +identity: + name: get_worksheet_pivot_data + author: Ryan Tian + label: + en_US: Get Worksheet Pivot Data + zh_Hans: 获取工作表统计透视数据 +description: + human: + en_US: Retrieve statistical pivot table data from a specified worksheet + zh_Hans: 从指定的工作表中检索统计透视表数据 + llm: A tool for extracting statistical pivot table data from a specific worksheet, providing summarized information for analysis and reporting purposes. +parameters: + - name: appkey + type: secret-input + required: true + label: + en_US: App Key + zh_Hans: App Key + human_description: + en_US: The AppKey parameter for the HAP application, typically found in the application's API documentation. + zh_Hans: HAP 应用的 AppKey 参数,可以从应用 API 文档中查找到 + llm_description: the AppKey parameter for the HAP application + form: form + + - name: sign + type: secret-input + required: true + label: + en_US: Sign + zh_Hans: Sign + human_description: + en_US: The Sign parameter for the HAP application + zh_Hans: HAP 应用的 Sign 参数 + llm_description: the Sign parameter for the HAP application + form: form + + - name: worksheet_id + type: string + required: true + label: + en_US: Worksheet ID + zh_Hans: 工作表 ID + human_description: + en_US: The ID of the specified worksheet + zh_Hans: 要获取字段信息的工作表 ID + llm_description: The ID of the specified worksheet which to get the fields information. + form: llm + + - name: x_column_fields + type: string + required: true + label: + en_US: Columns (X-axis) + zh_Hans: 统计列字段(X轴) + human_description: + en_US: The column fields that make up the pivot table's X-axis groups or other dimensions for the X-axis in pivot charts + zh_Hans: 组成透视表的统计列或者统计图表的X轴分组及X轴其它维度。JSON 对象数组格式,数组元素属性:controlId-列ID,displayName-显示名称,particleSize(可选)-字段类型是日期或者地区时,通过此参数设置统计维度(日期时间:1-日,2-周,3-月;地区:1-全国,2-省,3-市) + llm_description: | + This parameter allows you to specify the columns that make up the pivot table's X-axis groups or other dimensions for the X-axis in pivot charts. It is formatted as a JSON array, with its structure defined as follows: + ``` + type XColumnFields = { // X-axis or column object array + controlId: string; // fieldId + displayName: string; // displayName + particleSize?: number; // field type is date or area, set the statistical dimension (date time: 1-day, 2-week, 3-month; area: 1-nation, 2-province, 3-city) + }[]; + ``` + form: llm + + - name: y_row_fields + type: string + required: false + label: + en_US: Rows (Y-axis) + zh_Hans: 统计行字段(Y轴) + human_description: + en_US: The row fields that make up the pivot table's Y-axis groups or other dimensions for the Y-axis in pivot charts + zh_Hans: 组成透视表的统计行或者统计图表的Y轴分组及Y轴其它维度。JSON 对象数组格式,数组元素属性:controlId-列ID,displayName-显示名称,particleSize(可选)-字段类型是日期或者地区时,通过此参数设置统计维度(日期时间:1-日,2-周,3-月;地区:1-全国,2-省,3-市) + llm_description: | + This parameter allows you to specify the rows that make up the pivot table's Y-axis groups or other dimensions for the Y-axis in pivot charts. It is formatted as a JSON array, with its structure defined as follows: + ``` + type YRowFields = { // Y-axis or row object array + controlId: string; // fieldId + displayName: string; // displayName + particleSize?: number; // field type is date or area, set the statistical dimension (date time: 1-day, 2-week, 3-month; area: 1-nation, 2-province, 3-city) + }[]; + ``` + form: llm + + - name: value_fields + type: string + required: true + label: + en_US: Aggregated Values + zh_Hans: 统计值字段 + human_description: + en_US: The aggregated value fields in the pivot table + zh_Hans: 透视表中经过聚合计算后的统计值字段。JSON 对象数组格式,数组元素属性:controlId-列ID,displayName-显示名称,aggregation-聚合方式(SUM,AVG,MIN,MAX,COUNT) + llm_description: | + This parameter allows you to specify the aggregated value fields in the pivot table. It is formatted as a JSON array, with its structure defined as follows: + ``` + type ValueFields = { // aggregated value object array + controlId: string; // fieldId + displayName: string; // displayName + aggregation: string; // aggregation method, e.g.: SUM, AVG, MIN, MAX, COUNT + }[]; + ``` + form: llm + + - name: filters + type: string + required: false + label: + en_US: Filter Set + zh_Hans: 筛选器组合 + human_description: + en_US: A combination of filters applied to query records, formatted as a JSON array. See the application's API documentation for details on its structure and usage. + zh_Hans: 查询记录的筛选条件组合,格式为 JSON 数组,可以从应用 API 文档中了解参数结构详情 + llm_description: | + This parameter allows you to specify a set of conditions that records must meet to be included in the result set. It is formatted as a JSON array, with its structure defined as follows: + ``` + type Filters = { // filter object array + controlId: string; // fieldId + dataType: number; // fieldTypeId + spliceType: number; // condition concatenation method, 1: And, 2: Or + filterType: number; // expression type, refer to the for enumerable values + values?: string[]; // values in the condition, for option-type fields, multiple values can be passed + value?: string; // value in the condition, a single value can be passed according to the field type + dateRange?: number; // date range, mandatory when filterType is 17 or 18, refer to the for enumerable values + minValue?: string; // minimum value for custom range + maxValue?: string; // maximum value for custom range + isAsc?: boolean; // ascending order, false: descending, true: ascending + }[]; + ``` + For option-type fields, if this option field has `options`, then you need to get the corresponding `key` value from the `options` in the current field information via `value`, and pass it into `values` in array format. Do not use the `options` value of other fields as input conditions. + + ### FilterTypeEnum Reference + ``` + Enum Value, Enum Character, Description + 1, Like, Contains + 2, Eq, Is (Equal) + 3, Start, Starts With + 4, End, Ends With + 5, NotLike, Does Not Contain + 6, Ne, Is Not (Not Equal) + 7, IsEmpty, Empty + 8, HasValue, Not Empty + 11, Between, Within Range + 12, NotBetween, Outside Range + 13, Gt, Greater Than + 14, Gte, Greater Than or Equal To + 15, Lt, Less Than + 16, Lte, Less Than or Equal To + 17, DateEnum, Date Is + 18, NotDateEnum, Date Is Not + 21, MySelf, Owned by Me + 22, UnRead, Unread + 23, Sub, Owned by Subordinate + 24, RCEq, Associated Field Is + 25, RCNe, Associated Field Is Not + 26, ArrEq, Array Equals + 27, ArrNe, Array Does Not Equal + 31, DateBetween, Date Within Range (can only be used with minValue and maxValue) + 32, DateNotBetween, Date Not Within Range (can only be used with minValue and maxValue) + 33, DateGt, Date Later Than + 34, DateGte, Date Later Than or Equal To + 35, DateLt, Date Earlier Than + 36, DateLte, Date Earlier Than or Equal To + ``` + + ### DateRangeEnum Reference + ``` + Enum Value, Enum Character, Description + 1, Today, Today + 2, Yesterday, Yesterday + 3, Tomorrow, Tomorrow + 4, ThisWeek, This Week + 5, LastWeek, Last Week + 6, NextWeek, Next Week + 7, ThisMonth, This Month + 8, LastMonth, Last Month + 9, NextMonth, Next Month + 12, ThisQuarter, This Quarter + 13, LastQuarter, Last Quarter + 14, NextQuarter, Next Quarter + 15, ThisYear, This Year + 16, LastYear, Last Year + 17, NextYear, Next Year + 18, Customize, Custom + 21, Last7Day, Past 7 Days + 22, Last14Day, Past 14 Days + 23, Last30Day, Past 30 Days + 31, Next7Day, Next 7 Days + 32, Next14Day, Next 14 Days + 33, Next33Day, Next 33 Days + ``` + form: llm + + - name: sort_fields + type: string + required: false + label: + en_US: Sort Fields + zh_Hans: 排序字段 + human_description: + en_US: The fields to used for sorting + zh_Hans: 用于确定排序的字段,不超过3个 + llm_description: | + This optional parameter specifies the unique identifier of the fields that will be used to sort the results. It is in the format of an array of JSON objects, and its structure is defined as follows: + ``` + type SortByFields = { + controlId: string; // Field ID used for sorting + isAsc: boolean; // Sorting direction, true indicates ascending order, false indicates descending order + }[]; + ``` + form: llm + + - name: host + type: string + required: false + label: + en_US: Host Address + zh_Hans: 服务器地址 + human_description: + en_US: The address for the privately deployed HAP server. + zh_Hans: 私有部署 HAP 服务器地址,公有云无需填写 + llm_description: the address for the privately deployed HAP server. + form: form + + - name: result_type + type: select + required: true + options: + - value: table + label: + en_US: table text + zh_Hans: 表格文本 + - value: json + label: + en_US: json text + zh_Hans: JSON文本 + default: table + label: + en_US: Result type + zh_Hans: 结果类型 + human_description: + en_US: used for selecting the result type, table styled text or json text + zh_Hans: 用于选择结果类型,使用表格格式文本还是JSON格式文本 + form: form diff --git a/api/core/tools/provider/builtin/hap/tools/list_worksheet_records.py b/api/core/tools/provider/builtin/hap/tools/list_worksheet_records.py new file mode 100644 index 0000000000000000000000000000000000000000..9e43d5c532e5b377b9b6a3dad2b4637dd01ca1c3 --- /dev/null +++ b/api/core/tools/provider/builtin/hap/tools/list_worksheet_records.py @@ -0,0 +1,231 @@ +import json +import re +from typing import Any, Union + +import httpx + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class ListWorksheetRecordsTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + appkey = tool_parameters.get("appkey", "") + if not appkey: + return self.create_text_message("Invalid parameter App Key") + + sign = tool_parameters.get("sign", "") + if not sign: + return self.create_text_message("Invalid parameter Sign") + + worksheet_id = tool_parameters.get("worksheet_id", "") + if not worksheet_id: + return self.create_text_message("Invalid parameter Worksheet ID") + + host = tool_parameters.get("host", "") + if not host: + host = "https://api.mingdao.com" + elif not (host.startswith("http://") or host.startswith("https://")): + return self.create_text_message("Invalid parameter Host Address") + else: + host = f"{host.removesuffix('/')}/api" + + url_fields = f"{host}/v2/open/worksheet/getWorksheetInfo" + headers = {"Content-Type": "application/json"} + payload = {"appKey": appkey, "sign": sign, "worksheetId": worksheet_id} + + field_ids = tool_parameters.get("field_ids", "") + + try: + res = httpx.post(url_fields, headers=headers, json=payload, timeout=30) + res_json = res.json() + if res.is_success: + if res_json["error_code"] != 1: + return self.create_text_message( + "Failed to get the worksheet information. {}".format(res_json["error_msg"]) + ) + else: + worksheet_name = res_json["data"]["name"] + fields, schema, table_header = self.get_schema(res_json["data"]["controls"], field_ids) + else: + return self.create_text_message( + f"Failed to get the worksheet information, status code: {res.status_code}, response: {res.text}" + ) + except Exception as e: + return self.create_text_message( + "Failed to get the worksheet information, something went wrong: {}".format(e) + ) + + if field_ids: + payload["controls"] = [v.strip() for v in field_ids.split(",")] if field_ids else [] + filters = tool_parameters.get("filters", "") + if filters: + payload["filters"] = json.loads(filters) + sort_id = tool_parameters.get("sort_id", "") + sort_is_asc = tool_parameters.get("sort_is_asc", False) + if sort_id: + payload["sortId"] = sort_id + payload["isAsc"] = sort_is_asc + limit = tool_parameters.get("limit", 50) + payload["pageSize"] = limit + page_index = tool_parameters.get("page_index", 1) + payload["pageIndex"] = page_index + payload["useControlId"] = True + payload["listType"] = 1 + + url = f"{host}/v2/open/worksheet/getFilterRows" + try: + res = httpx.post(url, headers=headers, json=payload, timeout=90) + res_json = res.json() + if res.is_success: + if res_json["error_code"] != 1: + return self.create_text_message("Failed to get the records. {}".format(res_json["error_msg"])) + else: + result = { + "fields": fields, + "rows": [], + "total": res_json.get("data", {}).get("total"), + "payload": { + key: payload[key] + for key in [ + "worksheetId", + "controls", + "filters", + "sortId", + "isAsc", + "pageSize", + "pageIndex", + ] + if key in payload + }, + } + rows = res_json.get("data", {}).get("rows", []) + result_type = tool_parameters.get("result_type", "") + if not result_type: + result_type = "table" + if result_type == "json": + for row in rows: + result["rows"].append(self.get_row_field_value(row, schema)) + return self.create_text_message(json.dumps(result, ensure_ascii=False)) + else: + result_text = f'Found {result["total"]} rows in worksheet "{worksheet_name}".' + if result["total"] > 0: + result_text += ( + f" The following are {min(limit, result['total'])}" + f" pieces of data presented in a table format:\n\n{table_header}" + ) + for row in rows: + result_values = [] + for f in fields: + result_values.append( + self.handle_value_type(row[f["fieldId"]], schema[f["fieldId"]]) + ) + result_text += "\n|" + "|".join(result_values) + "|" + return self.create_text_message(result_text) + else: + return self.create_text_message( + f"Failed to get the records, status code: {res.status_code}, response: {res.text}" + ) + except Exception as e: + return self.create_text_message("Failed to get the records, something went wrong: {}".format(e)) + + def get_row_field_value(self, row: dict, schema: dict): + row_value = {"rowid": row["rowid"]} + for field in schema: + row_value[field] = self.handle_value_type(row[field], schema[field]) + return row_value + + def get_schema(self, controls: list, fieldids: str): + allow_fields = {v.strip() for v in fieldids.split(",")} if fieldids else set() + fields = [] + schema = {} + field_names = [] + for control in controls: + control_type_id = self.get_real_type_id(control) + if (control_type_id in self._get_ignore_types()) or ( + allow_fields and control["controlId"] not in allow_fields + ): + continue + else: + fields.append({"fieldId": control["controlId"], "fieldName": control["controlName"]}) + schema[control["controlId"]] = {"typeId": control_type_id, "options": self.set_option(control)} + field_names.append(control["controlName"]) + if not allow_fields or ("ctime" in allow_fields): + fields.append({"fieldId": "ctime", "fieldName": "Created Time"}) + schema["ctime"] = {"typeId": 16, "options": {}} + field_names.append("Created Time") + fields.append({"fieldId": "rowid", "fieldName": "Record Row ID"}) + schema["rowid"] = {"typeId": 2, "options": {}} + field_names.append("Record Row ID") + return fields, schema, "|" + "|".join(field_names) + "|\n|" + "---|" * len(field_names) + + def get_real_type_id(self, control: dict) -> int: + return control["sourceControlType"] if control["type"] == 30 else control["type"] + + def set_option(self, control: dict) -> dict: + options = {} + if control.get("options"): + options = {option["key"]: option["value"] for option in control["options"]} + elif control.get("advancedSetting", {}).get("itemnames"): + try: + itemnames = json.loads(control["advancedSetting"]["itemnames"]) + options = {item["key"]: item["value"] for item in itemnames} + except json.JSONDecodeError: + pass + return options + + def _get_ignore_types(self): + return {14, 21, 22, 34, 42, 43, 45, 47, 49, 10010} + + def handle_value_type(self, value, field): + type_id = field.get("typeId") + if type_id == 10: + value = value if isinstance(value, str) else "、".join(value) + elif type_id in {28, 36}: + value = field.get("options", {}).get(value, value) + elif type_id in {26, 27, 48, 14}: + value = self.process_value(value) + elif type_id in {35, 29}: + value = self.parse_cascade_or_associated(field, value) + elif type_id == 40: + value = self.parse_location(value) + return self.rich_text_to_plain_text(value) if value else "" + + def process_value(self, value): + if isinstance(value, str): + if value.startswith('[{"accountId"'): + value = json.loads(value) + value = ", ".join([item["fullname"] for item in value]) + elif value.startswith('[{"departmentId"'): + value = json.loads(value) + value = "、".join([item["departmentName"] for item in value]) + elif value.startswith('[{"organizeId"'): + value = json.loads(value) + value = "、".join([item["organizeName"] for item in value]) + elif value.startswith('[{"file_id"') or value == "[]": + value = "" + elif hasattr(value, "accountId"): + value = value["fullname"] + return value + + def parse_cascade_or_associated(self, field, value): + if (field["typeId"] == 35 and value.startswith("[")) or (field["typeId"] == 29 and value.startswith("[{")): + value = json.loads(value) + value = value[0]["name"] if len(value) > 0 else "" + else: + value = "" + return value + + def parse_location(self, value): + if len(value) > 10: + parsed_value = json.loads(value) + value = parsed_value.get("address", "") + else: + value = "" + return value + + def rich_text_to_plain_text(self, rich_text): + text = re.sub(r"<[^>]+>", "", rich_text) if "<" in rich_text else rich_text + return text.replace("|", "▏").replace("\n", " ") diff --git a/api/core/tools/provider/builtin/hap/tools/list_worksheet_records.yaml b/api/core/tools/provider/builtin/hap/tools/list_worksheet_records.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3c37746b921d058b861281bebb267c744d7713a0 --- /dev/null +++ b/api/core/tools/provider/builtin/hap/tools/list_worksheet_records.yaml @@ -0,0 +1,226 @@ +identity: + name: list_worksheet_records + author: Ryan Tian + label: + en_US: List Worksheet Records + zh_Hans: 查询工作表记录数据 +description: + human: + en_US: List records from the worksheet + zh_Hans: 查询工作表的记录列表数据,一次最多1000行,可分页获取 + llm: A tool to retrieve record data from the specific worksheet. +parameters: + - name: appkey + type: secret-input + required: true + label: + en_US: App Key + zh_Hans: App Key + human_description: + en_US: The AppKey parameter for the HAP application, typically found in the application's API documentation. + zh_Hans: HAP 应用的 AppKey 参数,可以从应用 API 文档中查找到 + llm_description: the AppKey parameter for the HAP application + form: form + + - name: sign + type: secret-input + required: true + label: + en_US: Sign + zh_Hans: Sign + human_description: + en_US: The Sign parameter for the HAP application + zh_Hans: HAP 应用的 Sign 参数 + llm_description: the Sign parameter for the HAP application + form: form + + - name: worksheet_id + type: string + required: true + label: + en_US: Worksheet ID + zh_Hans: 工作表 ID + human_description: + en_US: The ID of the worksheet from which to retrieve record data + zh_Hans: 要获取记录数据的工作表 ID + llm_description: This parameter specifies the ID of the worksheet where the records are stored. + form: llm + + - name: field_ids + type: string + required: false + label: + en_US: Field IDs + zh_Hans: 字段 ID 列表 + human_description: + en_US: A comma-separated list of field IDs whose data to retrieve. If not provided, all fields' data will be fetched + zh_Hans: 要获取记录数据的字段 ID,多个 ID 间用英文逗号隔开,不传此参数则将获取所有字段的数据 + llm_description: This optional parameter lets you specify a comma-separated list of field IDs. Unless the user explicitly requests to output the specified field in the question, this parameter should usually be omitted. If this parameter is omitted, the API will return data for all fields by default. When provided, only the data associated with these fields will be included in the response. + form: llm + + - name: filters + type: string + required: false + label: + en_US: Filter Set + zh_Hans: 筛选器组合 + human_description: + en_US: A combination of filters applied to query records, formatted as a JSON array. See the application's API documentation for details on its structure and usage. + zh_Hans: 查询记录的筛选条件组合,格式为 JSON 数组,可以从应用 API 文档中了解参数结构详情 + llm_description: | + This parameter allows you to specify a set of conditions that records must meet to be included in the result set. It is formatted as a JSON array, with its structure defined as follows: + ``` + type Filters = { // filter object array + controlId: string; // fieldId + dataType: number; // fieldTypeId + spliceType: number; // condition concatenation method, 1: And, 2: Or + filterType: number; // expression type, refer to the for enumerable values + values?: string[]; // values in the condition, for option-type fields, multiple values can be passed + value?: string; // value in the condition, a single value can be passed according to the field type + dateRange?: number; // date range, mandatory when filterType is 17 or 18, refer to the for enumerable values + minValue?: string; // minimum value for custom range + maxValue?: string; // maximum value for custom range + isAsc?: boolean; // ascending order, false: descending, true: ascending + }[]; + ``` + For option-type fields, if this option field has `options`, then you need to get the corresponding `key` value from the `options` in the current field information via `value`, and pass it into `values` in array format. Do not use the `options` value of other fields as input conditions. + + ### FilterTypeEnum Reference + ``` + Enum Value, Enum Character, Description + 1, Like, Contains(Include) + 2, Eq, Is (Equal) + 3, Start, Starts With + 4, End, Ends With + 5, NotLike, Does Not Contain(Not Include) + 6, Ne, Is Not (Not Equal) + 7, IsEmpty, Empty + 8, HasValue, Not Empty + 11, Between, Within Range(Belong to) + 12, NotBetween, Outside Range(Not belong to) + 13, Gt, Greater Than + 14, Gte, Greater Than or Equal To + 15, Lt, Less Than + 16, Lte, Less Than or Equal To + 17, DateEnum, Date Is + 18, NotDateEnum, Date Is Not + 24, RCEq, Associated Field Is + 25, RCNe, Associated Field Is Not + 26, ArrEq, Array Equals + 27, ArrNe, Array Does Not Equal + 31, DateBetween, Date Within Range (can only be used with minValue and maxValue) + 32, DateNotBetween, Date Not Within Range (can only be used with minValue and maxValue) + 33, DateGt, Date Later Than + 34, DateGte, Date Later Than or Equal To + 35, DateLt, Date Earlier Than + 36, DateLte, Date Earlier Than or Equal To + ``` + + ### DateRangeEnum Reference + ``` + Enum Value, Enum Character, Description + 1, Today, Today + 2, Yesterday, Yesterday + 3, Tomorrow, Tomorrow + 4, ThisWeek, This Week + 5, LastWeek, Last Week + 6, NextWeek, Next Week + 7, ThisMonth, This Month + 8, LastMonth, Last Month + 9, NextMonth, Next Month + 12, ThisQuarter, This Quarter + 13, LastQuarter, Last Quarter + 14, NextQuarter, Next Quarter + 15, ThisYear, This Year + 16, LastYear, Last Year + 17, NextYear, Next Year + 18, Customize, Custom + 21, Last7Day, Past 7 Days + 22, Last14Day, Past 14 Days + 23, Last30Day, Past 30 Days + 31, Next7Day, Next 7 Days + 32, Next14Day, Next 14 Days + 33, Next33Day, Next 33 Days + ``` + form: llm + + - name: sort_id + type: string + required: false + label: + en_US: Sort Field ID + zh_Hans: 排序字段 ID + human_description: + en_US: The ID of the field used for sorting + zh_Hans: 用以排序的字段 ID + llm_description: This optional parameter specifies the unique identifier of the field that will be used to sort the results. It should be set to the ID of an existing field within your data structure. + form: llm + + - name: sort_is_asc + type: boolean + required: false + label: + en_US: Ascending Order + zh_Hans: 是否升序排列 + human_description: + en_US: Determines whether the sorting is in ascending (true) or descending (false) order + zh_Hans: 排序字段的排序方式:true-升序,false-降序 + llm_description: This optional parameter controls the direction of the sort. If set to true, the results will be sorted in ascending order; if false, they will be sorted in descending order. + form: llm + + - name: limit + type: number + required: false + label: + en_US: Record Limit + zh_Hans: 记录数量限制 + human_description: + en_US: The maximum number of records to retrieve + zh_Hans: 要获取的记录数量限制条数 + llm_description: This optional parameter allows you to specify the maximum number of records that should be returned in the result set. When retrieving paginated record data, this parameter indicates the number of rows to fetch per page, and must be used in conjunction with the `page_index` parameter. + form: llm + + - name: page_index + type: number + required: false + label: + en_US: Page Index + zh_Hans: 页码 + human_description: + en_US: The page number when paginating through a list of records + zh_Hans: 分页读取记录列表时的页码 + llm_description: This parameter is used when you need to paginate through a large set of records. The default value is 1, which refers to the first page. When it is used, the meaning of the `limit` parameter becomes the number of records per page. + form: llm + + - name: host + type: string + required: false + label: + en_US: Host Address + zh_Hans: 服务器地址 + human_description: + en_US: The address for the privately deployed HAP server. + zh_Hans: 私有部署 HAP 服务器地址,公有云无需填写 + llm_description: the address for the privately deployed HAP server. + form: form + + - name: result_type + type: select + required: true + options: + - value: table + label: + en_US: table text + zh_Hans: 表格文本 + - value: json + label: + en_US: json text + zh_Hans: JSON文本 + default: table + label: + en_US: Result type + zh_Hans: 结果类型 + human_description: + en_US: used for selecting the result type, table styled text or json text + zh_Hans: 用于选择结果类型,使用表格格式文本还是JSON格式文本 + form: form diff --git a/api/core/tools/provider/builtin/hap/tools/list_worksheets.py b/api/core/tools/provider/builtin/hap/tools/list_worksheets.py new file mode 100644 index 0000000000000000000000000000000000000000..4e852c0028497c0cabf029424c1f828238df10ab --- /dev/null +++ b/api/core/tools/provider/builtin/hap/tools/list_worksheets.py @@ -0,0 +1,83 @@ +import json +from typing import Any, Union + +import httpx + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class ListWorksheetsTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + appkey = tool_parameters.get("appkey", "") + if not appkey: + return self.create_text_message("Invalid parameter App Key") + sign = tool_parameters.get("sign", "") + if not sign: + return self.create_text_message("Invalid parameter Sign") + + host = tool_parameters.get("host", "") + if not host: + host = "https://api.mingdao.com" + elif not (host.startswith("http://") or host.startswith("https://")): + return self.create_text_message("Invalid parameter Host Address") + else: + host = f"{host.removesuffix('/')}/api" + url = f"{host}/v1/open/app/get" + + result_type = tool_parameters.get("result_type", "") + if not result_type: + result_type = "table" + + headers = {"Content-Type": "application/json"} + params = { + "appKey": appkey, + "sign": sign, + } + try: + res = httpx.get(url, headers=headers, params=params, timeout=30) + res_json = res.json() + if res.is_success: + if res_json["error_code"] != 1: + return self.create_text_message( + "Failed to access the application. {}".format(res_json["error_msg"]) + ) + else: + if result_type == "json": + worksheets = [] + for section in res_json["data"]["sections"]: + worksheets.extend(self._extract_worksheets(section, result_type)) + return self.create_text_message(text=json.dumps(worksheets, ensure_ascii=False)) + else: + worksheets = "|worksheetId|worksheetName|description|\n|---|---|---|" + for section in res_json["data"]["sections"]: + worksheets += self._extract_worksheets(section, result_type) + return self.create_text_message(worksheets) + + else: + return self.create_text_message( + f"Failed to list worksheets, status code: {res.status_code}, response: {res.text}" + ) + except Exception as e: + return self.create_text_message("Failed to list worksheets, something went wrong: {}".format(e)) + + def _extract_worksheets(self, section, type): + items = [] + tables = "" + for item in section.get("items", []): + if item.get("type") == 0 and ("notes" not in item or item.get("notes") != "NO"): + if type == "json": + filtered_item = {"id": item["id"], "name": item["name"], "notes": item.get("notes", "")} + items.append(filtered_item) + else: + tables += f"\n|{item['id']}|{item['name']}|{item.get('notes', '')}|" + + for child_section in section.get("childSections", []): + if type == "json": + items.extend(self._extract_worksheets(child_section, "json")) + else: + tables += self._extract_worksheets(child_section, "table") + + return items if type == "json" else tables diff --git a/api/core/tools/provider/builtin/hap/tools/list_worksheets.yaml b/api/core/tools/provider/builtin/hap/tools/list_worksheets.yaml new file mode 100644 index 0000000000000000000000000000000000000000..935b72a89564cd628b524a026f0d91139a8792cc --- /dev/null +++ b/api/core/tools/provider/builtin/hap/tools/list_worksheets.yaml @@ -0,0 +1,68 @@ +identity: + name: list_worksheets + author: Ryan Tian + label: + en_US: List Worksheets + zh_Hans: 获取应用下所有工作表 +description: + human: + en_US: List worksheets within an application + zh_Hans: 获取应用下的所有工作表和说明信息 + llm: A tool to list worksheets info within an application, imported parameter is AppKey and Sign of the application. +parameters: + - name: appkey + type: secret-input + required: true + label: + en_US: App Key + zh_Hans: App Key + human_description: + en_US: The AppKey parameter for the HAP application, typically found in the application's API documentation. + zh_Hans: HAP 应用的 AppKey 参数,可以从应用 API 文档中查找到 + llm_description: the AppKey parameter for the HAP application + form: form + + - name: sign + type: secret-input + required: true + label: + en_US: Sign + zh_Hans: Sign + human_description: + en_US: The Sign parameter for the HAP application + zh_Hans: HAP 应用的 Sign 参数 + llm_description: the Sign parameter for the HAP application + form: form + + - name: host + type: string + required: false + label: + en_US: Host Address + zh_Hans: 服务器地址 + human_description: + en_US: The address for the privately deployed HAP server. + zh_Hans: 私有部署 HAP 服务器地址,公有云无需填写 + llm_description: the address for the privately deployed HAP server. + form: form + + - name: result_type + type: select + required: true + options: + - value: table + label: + en_US: table text + zh_Hans: 表格文本 + - value: json + label: + en_US: json text + zh_Hans: JSON文本 + default: table + label: + en_US: Result type + zh_Hans: 结果类型 + human_description: + en_US: used for selecting the result type, table styled text or json text + zh_Hans: 用于选择结果类型,使用表格格式文本还是JSON格式文本 + form: form diff --git a/api/core/tools/provider/builtin/hap/tools/update_worksheet_record.py b/api/core/tools/provider/builtin/hap/tools/update_worksheet_record.py new file mode 100644 index 0000000000000000000000000000000000000000..971f3d37f6dfbf0744c2748fb86357110140f1a1 --- /dev/null +++ b/api/core/tools/provider/builtin/hap/tools/update_worksheet_record.py @@ -0,0 +1,55 @@ +import json +from typing import Any, Union + +import httpx + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class UpdateWorksheetRecordTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + appkey = tool_parameters.get("appkey", "") + if not appkey: + return self.create_text_message("Invalid parameter App Key") + sign = tool_parameters.get("sign", "") + if not sign: + return self.create_text_message("Invalid parameter Sign") + worksheet_id = tool_parameters.get("worksheet_id", "") + if not worksheet_id: + return self.create_text_message("Invalid parameter Worksheet ID") + row_id = tool_parameters.get("row_id", "") + if not row_id: + return self.create_text_message("Invalid parameter Record Row ID") + record_data = tool_parameters.get("record_data", "") + if not record_data: + return self.create_text_message("Invalid parameter Record Row Data") + + host = tool_parameters.get("host", "") + if not host: + host = "https://api.mingdao.com" + elif not host.startswith(("http://", "https://")): + return self.create_text_message("Invalid parameter Host Address") + else: + host = f"{host.removesuffix('/')}/api" + + url = f"{host}/v2/open/worksheet/editRow" + headers = {"Content-Type": "application/json"} + payload = {"appKey": appkey, "sign": sign, "worksheetId": worksheet_id, "rowId": row_id} + + try: + payload["controls"] = json.loads(record_data) + res = httpx.post(url, headers=headers, json=payload, timeout=60) + res.raise_for_status() + res_json = res.json() + if res_json.get("error_code") != 1: + return self.create_text_message(f"Failed to update the record. {res_json['error_msg']}") + return self.create_text_message("Record updated successfully.") + except httpx.RequestError as e: + return self.create_text_message(f"Failed to update the record, request error: {e}") + except json.JSONDecodeError as e: + return self.create_text_message(f"Failed to parse JSON response: {e}") + except Exception as e: + return self.create_text_message(f"Failed to update the record, unexpected error: {e}") diff --git a/api/core/tools/provider/builtin/hap/tools/update_worksheet_record.yaml b/api/core/tools/provider/builtin/hap/tools/update_worksheet_record.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fe1f8f671a4e2f084ba24c7ca21b119c9918b929 --- /dev/null +++ b/api/core/tools/provider/builtin/hap/tools/update_worksheet_record.yaml @@ -0,0 +1,90 @@ +identity: + name: update_worksheet_record + author: Ryan Tian + label: + en_US: Update Worksheet Record + zh_Hans: 更新指定的一条工作表记录 +description: + human: + en_US: Updates a single record in a worksheet based on the specified record row ID + zh_Hans: 根据指定的记录ID更新一条工作表记录数据 + llm: A tool to modify existing information within a particular record of a worksheet by referencing its unique identifier. +parameters: + - name: appkey + type: secret-input + required: true + label: + en_US: App Key + zh_Hans: App Key + human_description: + en_US: The AppKey parameter for the HAP application, typically found in the application's API documentation. + zh_Hans: HAP 应用的 AppKey 参数,可以从应用 API 文档中查找到 + llm_description: the AppKey parameter for the HAP application + form: form + + - name: sign + type: secret-input + required: true + label: + en_US: Sign + zh_Hans: Sign + human_description: + en_US: The Sign parameter for the HAP application + zh_Hans: HAP 应用的 Sign 参数 + llm_description: the Sign parameter for the HAP application + form: form + + - name: worksheet_id + type: string + required: true + label: + en_US: Worksheet ID + zh_Hans: 工作表 ID + human_description: + en_US: The ID of the specified worksheet + zh_Hans: 要获取字段信息的工作表 ID + llm_description: The ID of the specified worksheet which to get the fields information. + form: llm + + - name: row_id + type: string + required: true + label: + en_US: Record Row ID + zh_Hans: 记录 ID + human_description: + en_US: The row ID of the specified record + zh_Hans: 要更新的记录 ID + llm_description: The row ID of the specified record which to be updated. + form: llm + + - name: record_data + type: string + required: true + label: + en_US: Record Row Data + zh_Hans: 记录数据 + human_description: + en_US: The fields with data of the specified record + zh_Hans: 要更新的记录数据,JSON 对象数组格式。数组元素属性:controlId-字段ID,value-字段值 + llm_description: | + The fields with data of the specified record which to be updated. It is in the format of an array of JSON objects, and the structure is defined as follows: + ``` + type RowData = { + controlId: string; // Field ID to be updated + value: string; // Field value to be updated + }[]; + ``` + form: llm + + - name: host + type: string + required: false + label: + en_US: Host Address + zh_Hans: 服务器地址 + human_description: + en_US: The address for the privately deployed HAP server. + zh_Hans: 私有部署 HAP 服务器地址,公有云无需填写 + llm_description: the address for the privately deployed HAP server. + form: form diff --git a/api/core/tools/provider/builtin/jina/_assets/icon.svg b/api/core/tools/provider/builtin/jina/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..2e1b00fa52e43c7affeda36522bc22d0ab17d9a0 --- /dev/null +++ b/api/core/tools/provider/builtin/jina/_assets/icon.svg @@ -0,0 +1,4 @@ + + + + diff --git a/api/core/tools/provider/builtin/jina/jina.py b/api/core/tools/provider/builtin/jina/jina.py new file mode 100644 index 0000000000000000000000000000000000000000..154e15db016dd1897d047dfc886a6a3aee01896c --- /dev/null +++ b/api/core/tools/provider/builtin/jina/jina.py @@ -0,0 +1,38 @@ +import json +from typing import Any + +from core.tools.entities.values import ToolLabelEnum +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.jina.tools.jina_reader import JinaReaderTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class GoogleProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + if credentials["api_key"] is None: + credentials["api_key"] = "" + else: + result = ( + JinaReaderTool() + .fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ) + .invoke( + user_id="", + tool_parameters={ + "url": "https://example.com", + }, + )[0] + ) + + message = json.loads(result.message) + if message["code"] != 200: + raise ToolProviderCredentialValidationError(message["message"]) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) + + def _get_tool_labels(self) -> list[ToolLabelEnum]: + return [ToolLabelEnum.SEARCH, ToolLabelEnum.PRODUCTIVITY] diff --git a/api/core/tools/provider/builtin/jina/jina.yaml b/api/core/tools/provider/builtin/jina/jina.yaml new file mode 100644 index 0000000000000000000000000000000000000000..af3ca23ffaff464a9875f7af2b09862a00dbd0f9 --- /dev/null +++ b/api/core/tools/provider/builtin/jina/jina.yaml @@ -0,0 +1,32 @@ +identity: + author: Dify + name: jina + label: + en_US: Jina AI + zh_Hans: Jina AI + pt_BR: Jina AI + description: + en_US: Your Search Foundation, Supercharged! + zh_Hans: 您的搜索底座,从此不同! + pt_BR: Your Search Foundation, Supercharged! + icon: icon.svg + tags: + - search + - productivity +credentials_for_provider: + api_key: + type: secret-input + required: false + label: + en_US: API Key (leave empty if you don't have one) + zh_Hans: API 密钥(可留空) + pt_BR: Chave API (deixe vazio se você não tiver uma) + placeholder: + en_US: Please enter your Jina AI API key + zh_Hans: 请输入你的 Jina AI API 密钥 + pt_BR: Por favor, insira sua chave de API do Jina AI + help: + en_US: Get your Jina AI API key from Jina AI (optional, but you can get a higher rate) + zh_Hans: 从 Jina AI 获取您的 Jina AI API 密钥(非必须,能得到更高的速率) + pt_BR: Obtenha sua chave de API do Jina AI na Jina AI (opcional, mas você pode obter uma taxa mais alta) + url: https://jina.ai diff --git a/api/core/tools/provider/builtin/jina/tools/jina_reader.py b/api/core/tools/provider/builtin/jina/tools/jina_reader.py new file mode 100644 index 0000000000000000000000000000000000000000..756b72722481461cb5a0c53b0d14bd4d2819444c --- /dev/null +++ b/api/core/tools/provider/builtin/jina/tools/jina_reader.py @@ -0,0 +1,87 @@ +import json +from typing import Any, Union + +from yarl import URL + +from core.helper import ssrf_proxy +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class JinaReaderTool(BuiltinTool): + _jina_reader_endpoint = "https://r.jina.ai/" + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + url = tool_parameters["url"] + + headers = {"Accept": "application/json"} + + if "api_key" in self.runtime.credentials and self.runtime.credentials.get("api_key"): + headers["Authorization"] = "Bearer " + self.runtime.credentials.get("api_key") + + request_params = tool_parameters.get("request_params") + if request_params is not None and request_params != "": + try: + request_params = json.loads(request_params) + if not isinstance(request_params, dict): + raise ValueError("request_params must be a JSON object") + except (json.JSONDecodeError, ValueError) as e: + raise ValueError(f"Invalid request_params: {e}") + + target_selector = tool_parameters.get("target_selector") + if target_selector is not None and target_selector != "": + headers["X-Target-Selector"] = target_selector + + wait_for_selector = tool_parameters.get("wait_for_selector") + if wait_for_selector is not None and wait_for_selector != "": + headers["X-Wait-For-Selector"] = wait_for_selector + + remove_selector = tool_parameters.get("remove_selector") + if remove_selector is not None and remove_selector != "": + headers["X-Remove-Selector"] = remove_selector + + if tool_parameters.get("retain_images", False): + headers["X-Retain-Images"] = "true" + + if tool_parameters.get("image_caption", False): + headers["X-With-Generated-Alt"] = "true" + + if tool_parameters.get("gather_all_links_at_the_end", False): + headers["X-With-Links-Summary"] = "true" + + if tool_parameters.get("gather_all_images_at_the_end", False): + headers["X-With-Images-Summary"] = "true" + + proxy_server = tool_parameters.get("proxy_server") + if proxy_server is not None and proxy_server != "": + headers["X-Proxy-Url"] = proxy_server + + if tool_parameters.get("no_cache", False): + headers["X-No-Cache"] = "true" + + if tool_parameters.get("with_iframe", False): + headers["X-With-Iframe"] = "true" + + if tool_parameters.get("with_shadow_dom", False): + headers["X-With-Shadow-Dom"] = "true" + + max_retries = tool_parameters.get("max_retries", 3) + response = ssrf_proxy.get( + str(URL(self._jina_reader_endpoint + url)), + headers=headers, + params=request_params, + timeout=(10, 60), + max_retries=max_retries, + ) + + if tool_parameters.get("summary", False): + return self.create_text_message(self.summary(user_id, response.text)) + + return self.create_text_message(response.text) diff --git a/api/core/tools/provider/builtin/jina/tools/jina_reader.yaml b/api/core/tools/provider/builtin/jina/tools/jina_reader.yaml new file mode 100644 index 0000000000000000000000000000000000000000..012a8c7688cb57f8e6d2d36a1d442684aa94eaef --- /dev/null +++ b/api/core/tools/provider/builtin/jina/tools/jina_reader.yaml @@ -0,0 +1,221 @@ +identity: + name: jina_reader + author: Dify + label: + en_US: Fetch Single Page + zh_Hans: 获取单页面 + pt_BR: Fetch Single Page +description: + human: + en_US: Fetch the target URL (can be a PDF) and convert it into a LLM-friendly markdown. + zh_Hans: 获取目标网址(可以是 PDF),并将其转换为适合大模型处理的 Markdown 格式。 + pt_BR: Busque a URL de destino (que pode ser um PDF) e converta em um Markdown LLM-friendly. + llm: A tool for scraping webpages. Input should be a URL. +parameters: + - name: url + type: string + required: true + label: + en_US: URL + zh_Hans: 网址 + pt_BR: URL + human_description: + en_US: Web link + zh_Hans: 网页链接 + pt_BR: URL da web + llm_description: url para scraping + form: llm + - name: request_params + type: string + required: false + label: + en_US: Request params + zh_Hans: 请求参数 + pt_BR: Parâmetros de solicitação + human_description: + en_US: | + request parameters, format: {"key1": "value1", "key2": "value2"} + zh_Hans: | + 请求参数,格式:{"key1": "value1", "key2": "value2"} + pt_BR: | + parâmetros de solicitação, formato: {"key1": "value1", "key2": "value2"} + llm_description: request parameters + form: llm + - name: target_selector + type: string + required: false + label: + en_US: Target selector + zh_Hans: 目标选择器 + pt_BR: Seletor de destino + human_description: + en_US: css selector for scraping specific elements + zh_Hans: css 选择器用于抓取特定元素 + pt_BR: css selector para scraping de elementos específicos + llm_description: css selector of the target element to scrape + form: form + - name: wait_for_selector + type: string + required: false + label: + en_US: Wait for selector + zh_Hans: 等待选择器 + pt_BR: Aguardar por seletor + human_description: + en_US: css selector for waiting for specific elements + zh_Hans: css 选择器用于等待特定元素 + pt_BR: css selector para aguardar elementos específicos + llm_description: css selector of the target element to wait for + form: form + - name: remove_selector + type: string + required: false + label: + en_US: Excluded Selector + zh_Hans: 排除选择器 + pt_BR: Seletor Excluído + human_description: + en_US: css selector for remove for specific elements + zh_Hans: css 选择器用于排除特定元素 + pt_BR: seletor CSS para remover elementos específicos + llm_description: css selector of the target element to remove for + form: form + - name: retain_images + type: boolean + required: false + default: false + label: + en_US: Remove All Images + zh_Hans: 删除所有图片 + pt_BR: Remover todas as imagens + human_description: + en_US: Removes all images from the response. + zh_Hans: 从响应中删除所有图片。 + pt_BR: Remove todas as imagens da resposta. + llm_description: Remove all images + form: form + - name: image_caption + type: boolean + required: false + default: false + label: + en_US: Image caption + zh_Hans: 图片说明 + pt_BR: Legenda da imagem + human_description: + en_US: "Captions all images at the specified URL, adding 'Image [idx]: [caption]' as an alt tag for those without one. This allows downstream LLMs to interact with the images in activities such as reasoning and summarizing." + zh_Hans: "为指定 URL 上的所有图像添加标题,为没有标题的图像添加“Image [idx]: [caption]”作为 alt 标签,以支持下游模型的图像交互。" + pt_BR: "Adiciona legendas a todas as imagens na URL especificada, adicionando 'Imagem [idx]: [legenda]' como uma tag alt para aquelas que não têm uma. Isso permite que os modelos LLM inferiores interajam com as imagens em atividades como raciocínio e resumo." + llm_description: Captions all images at the specified URL + form: form + - name: gather_all_links_at_the_end + type: boolean + required: false + default: false + label: + en_US: Gather all links at the end + zh_Hans: 将所有链接集中到最后 + pt_BR: Coletar todos os links ao final + human_description: + en_US: A "Buttons & Links" section will be created at the end. This helps the downstream LLMs or web agents navigating the page or take further actions. + zh_Hans: 末尾将添加“按钮和链接”部分,方便下游模型或网络代理做页面导航或执行进一步操作。 + pt_BR: Um "Botões & Links" section will be created at the end. This helps the downstream LLMs or web agents navigating the page or take further actions. + llm_description: Gather all links at the end + form: form + - name: gather_all_images_at_the_end + type: boolean + required: false + default: false + label: + en_US: Gather all images at the end + zh_Hans: 将所有图片集中到最后 + pt_BR: Coletar todas as imagens ao final + human_description: + en_US: An "Images" section will be created at the end. This gives the downstream LLMs an overview of all visuals on the page, which may improve reasoning. + zh_Hans: 末尾会新增“图片”部分,方便下游模型全面了解页面的视觉内容,提升推理效果。 + pt_BR: Um "Imagens" section will be created at the end. This gives the downstream LLMs an overview of all visuals on the page, which may improve reasoning. + llm_description: Gather all images at the end + form: form + - name: proxy_server + type: string + required: false + label: + en_US: Proxy server + zh_Hans: 代理服务器 + pt_BR: Servidor de proxy + human_description: + en_US: Use proxy to access URLs + zh_Hans: 利用代理访问 URL + pt_BR: Use proxy to access URLs + llm_description: Use proxy to access URLs + form: form + - name: no_cache + type: boolean + required: false + default: false + label: + en_US: Bypass the Cache + zh_Hans: 绕过缓存 + pt_BR: Ignorar o cache + human_description: + en_US: Bypass the Cache + zh_Hans: 是否绕过缓存 + pt_BR: Ignorar o cache + llm_description: bypass the cache + form: form + - name: with_iframe + type: boolean + required: false + default: false + label: + en_US: Enable iframe extraction + zh_Hans: 启用 iframe 提取 + pt_BR: Habilitar extração de iframe + human_description: + en_US: Extract and process content of all embedded iframes in the DOM tree. + zh_Hans: 提取并处理 DOM 树中所有嵌入 iframe 的内容。 + pt_BR: Extrair e processar o conteúdo de todos os iframes incorporados na árvore DOM. + llm_description: Extract content from embedded iframes + form: form + - name: with_shadow_dom + type: boolean + required: false + default: false + label: + en_US: Enable Shadow DOM extraction + zh_Hans: 启用 Shadow DOM 提取 + pt_BR: Habilitar extração de Shadow DOM + human_description: + en_US: Traverse all Shadow DOM roots in the document and extract content. + zh_Hans: 遍历文档中所有 Shadow DOM 根并提取内容。 + pt_BR: Percorra todas as raízes do Shadow DOM no documento e extraia o conteúdo. + llm_description: Extract content from Shadow DOM roots + form: form + - name: summary + type: boolean + required: false + default: false + label: + en_US: Enable summary + zh_Hans: 是否启用摘要 + pt_BR: Habilitar resumo + human_description: + en_US: Enable summary for the output + zh_Hans: 为输出启用摘要 + pt_BR: Habilitar resumo para a saída + llm_description: enable summary + form: form + - name: max_retries + type: number + required: false + default: 3 + label: + en_US: Retry + zh_Hans: 重试 + pt_BR: Repetir + human_description: + en_US: Number of times to retry the request if it fails + zh_Hans: 请求失败时重试的次数 + pt_BR: Número de vezes para repetir a solicitação se falhar + llm_description: Number of times to retry the request if it fails + form: form diff --git a/api/core/tools/provider/builtin/jina/tools/jina_search.py b/api/core/tools/provider/builtin/jina/tools/jina_search.py new file mode 100644 index 0000000000000000000000000000000000000000..30af6de7831e590460110d15fb102a1f73c7ae42 --- /dev/null +++ b/api/core/tools/provider/builtin/jina/tools/jina_search.py @@ -0,0 +1,46 @@ +from typing import Any, Union + +from yarl import URL + +from core.helper import ssrf_proxy +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class JinaSearchTool(BuiltinTool): + _jina_search_endpoint = "https://s.jina.ai/" + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + query = tool_parameters["query"] + + headers = {"Accept": "application/json"} + + if "api_key" in self.runtime.credentials and self.runtime.credentials.get("api_key"): + headers["Authorization"] = "Bearer " + self.runtime.credentials.get("api_key") + + if tool_parameters.get("image_caption", False): + headers["X-With-Generated-Alt"] = "true" + + if tool_parameters.get("gather_all_links_at_the_end", False): + headers["X-With-Links-Summary"] = "true" + + if tool_parameters.get("gather_all_images_at_the_end", False): + headers["X-With-Images-Summary"] = "true" + + proxy_server = tool_parameters.get("proxy_server") + if proxy_server is not None and proxy_server != "": + headers["X-Proxy-Url"] = proxy_server + + if tool_parameters.get("no_cache", False): + headers["X-No-Cache"] = "true" + + max_retries = tool_parameters.get("max_retries", 3) + response = ssrf_proxy.get( + str(URL(self._jina_search_endpoint + query)), headers=headers, timeout=(10, 60), max_retries=max_retries + ) + + return self.create_text_message(response.text) diff --git a/api/core/tools/provider/builtin/jina/tools/jina_search.yaml b/api/core/tools/provider/builtin/jina/tools/jina_search.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e58c639e5690d096d79ec6faade406e34cf738fe --- /dev/null +++ b/api/core/tools/provider/builtin/jina/tools/jina_search.yaml @@ -0,0 +1,110 @@ +identity: + name: jina_search + author: Dify + label: + en_US: Search the web + zh_Hans: 联网搜索 + pt_BR: Search the web +description: + human: + en_US: Search on the public web of a given query and return the top results as LLM-friendly markdown. + zh_Hans: 针对给定的查询在互联网上进行搜索,并以适合大模型处理的 Markdown 格式返回最相关的结果。 + pt_BR: Procurar na web pública de uma consulta fornecida e retornar os melhores resultados como markdown para LLMs. + llm: A tool for searching results on the web for grounding. Input should be a simple question. +parameters: + - name: query + type: string + required: true + label: + en_US: Question (Query) + zh_Hans: 查询 + pt_BR: Pergunta (Consulta) + human_description: + en_US: used to find information on the web + zh_Hans: 在网络上搜索信息 + pt_BR: Usado para encontrar informações na web + llm_description: Pergunta simples para fazer na web + form: llm + - name: image_caption + type: boolean + required: false + default: false + label: + en_US: Image caption + zh_Hans: 图片说明 + pt_BR: Legenda da imagem + human_description: + en_US: "Captions all images at the specified URL, adding 'Image [idx]: [caption]' as an alt tag for those without one. This allows downstream LLMs to interact with the images in activities such as reasoning and summarizing." + zh_Hans: "为指定 URL 上的所有图像添加标题,为没有标题的图像添加“Image [idx]: [caption]”作为 alt 标签,以支持下游模型的图像交互。" + pt_BR: "Captions all images at the specified URL, adding 'Image [idx]: [caption]' as an alt tag for those without one. This allows downstream LLMs to interact with the images in activities such as reasoning and summarizing." + llm_description: Captions all images at the specified URL + form: form + - name: gather_all_links_at_the_end + type: boolean + required: false + default: false + label: + en_US: Gather all links at the end + zh_Hans: 将所有链接集中到最后 + pt_BR: Coletar todos os links ao final + human_description: + en_US: A "Buttons & Links" section will be created at the end. This helps the downstream LLMs or web agents navigating the page or take further actions. + zh_Hans: 末尾将添加“按钮和链接”部分,汇总页面上的所有链接。方便下游模型或网络代理做页面导航或执行进一步操作。 + pt_BR: Um "Botão & Links" seção será criada no final. Isso ajuda os LLMs ou agentes da web navegando pela página ou executar ações adicionais. + llm_description: Gather all links at the end + form: form + - name: gather_all_images_at_the_end + type: boolean + required: false + default: false + label: + en_US: Gather all images at the end + zh_Hans: 将所有图片集中到最后 + pt_BR: Coletar todas as imagens ao final + human_description: + en_US: An "Images" section will be created at the end. This gives the downstream LLMs an overview of all visuals on the page, which may improve reasoning. + zh_Hans: 末尾会新增“图片”部分,汇总页面上的所有图片。方便下游模型概览页面的视觉内容,提升推理效果。 + pt_BR: Um "Imagens" seção será criada no final. Isso fornece uma visão geral de todas as imagens na página para os LLMs, que pode melhorar a razão. + llm_description: Gather all images at the end + form: form + - name: proxy_server + type: string + required: false + label: + en_US: Proxy server + zh_Hans: 代理服务器 + pt_BR: Servidor de proxy + human_description: + en_US: Use proxy to access URLs + zh_Hans: 利用代理访问 URL + pt_BR: Usar proxy para acessar URLs + llm_description: Use proxy to access URLs + form: form + - name: no_cache + type: boolean + required: false + default: false + label: + en_US: Bypass the Cache + zh_Hans: 是否绕过缓存 + pt_BR: Ignorar o cache + human_description: + en_US: Bypass the Cache + zh_Hans: 是否绕过缓存 + pt_BR: Ignorar o cache + llm_description: bypass the cache + form: form + - name: max_retries + type: number + required: false + default: 3 + label: + en_US: Retry + zh_Hans: 重试 + pt_BR: Repetir + human_description: + en_US: Number of times to retry the request if it fails + zh_Hans: 请求失败时重试的次数 + pt_BR: Número de vezes para repetir a solicitação se falhar + llm_description: Number of times to retry the request if it fails + form: form diff --git a/api/core/tools/provider/builtin/jina/tools/jina_tokenizer.py b/api/core/tools/provider/builtin/jina/tools/jina_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..06dabcc9c2a74e8ce60cfb781d9a26bf2234a13d --- /dev/null +++ b/api/core/tools/provider/builtin/jina/tools/jina_tokenizer.py @@ -0,0 +1,39 @@ +from typing import Any + +from core.helper import ssrf_proxy +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class JinaTokenizerTool(BuiltinTool): + _jina_tokenizer_endpoint = "https://tokenize.jina.ai/" + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> ToolInvokeMessage: + content = tool_parameters["content"] + body = {"content": content} + + headers = {"Content-Type": "application/json"} + + if "api_key" in self.runtime.credentials and self.runtime.credentials.get("api_key"): + headers["Authorization"] = "Bearer " + self.runtime.credentials.get("api_key") + + if tool_parameters.get("return_chunks", False): + body["return_chunks"] = True + + if tool_parameters.get("return_tokens", False): + body["return_tokens"] = True + + if tokenizer := tool_parameters.get("tokenizer"): + body["tokenizer"] = tokenizer + + response = ssrf_proxy.post( + self._jina_tokenizer_endpoint, + headers=headers, + json=body, + ) + + return self.create_json_message(response.json()) diff --git a/api/core/tools/provider/builtin/jina/tools/jina_tokenizer.yaml b/api/core/tools/provider/builtin/jina/tools/jina_tokenizer.yaml new file mode 100644 index 0000000000000000000000000000000000000000..74885cdf9a70480f9ff5669423f001843b34ecf4 --- /dev/null +++ b/api/core/tools/provider/builtin/jina/tools/jina_tokenizer.yaml @@ -0,0 +1,78 @@ +identity: + name: jina_tokenizer + author: hjlarry + label: + en_US: Segment + zh_Hans: 切分器 + pt_BR: Segment +description: + human: + en_US: Split long text into chunks and do tokenization. + zh_Hans: 将长文本拆分成小段落,并做分词处理。 + pt_BR: Dividir o texto longo em pedaços e fazer tokenização. + llm: Free API to tokenize text and segment long text into chunks. +parameters: + - name: content + type: string + required: true + label: + en_US: Content + zh_Hans: 内容 + pt_BR: Conteúdo + llm_description: the content which need to tokenize or segment + form: llm + - name: return_tokens + type: boolean + required: false + label: + en_US: Return the tokens + zh_Hans: 是否返回tokens + pt_BR: Retornar os tokens + human_description: + en_US: Return the tokens and their corresponding ids in the response. + zh_Hans: 返回tokens及其对应的ids。 + pt_BR: Retornar os tokens e seus respectivos ids na resposta. + form: form + - name: return_chunks + type: boolean + label: + en_US: Return the chunks + zh_Hans: 是否分块 + pt_BR: Retornar os chunks + human_description: + en_US: Chunking the input into semantically meaningful segments while handling a wide variety of text types and edge cases based on common structural cues. + zh_Hans: 将输入文本分块为语义有意义的片段,同时基于常见的结构线索处理各种文本类型和特殊情况。 + pt_BR: Dividir o texto de entrada em segmentos semanticamente significativos, enquanto lida com uma ampla variedade de tipos de texto e casos de borda com base em pistas estruturais comuns. + form: form + - name: tokenizer + type: select + options: + - value: cl100k_base + label: + en_US: cl100k_base + - value: o200k_base + label: + en_US: o200k_base + - value: p50k_base + label: + en_US: p50k_base + - value: r50k_base + label: + en_US: r50k_base + - value: p50k_edit + label: + en_US: p50k_edit + - value: gpt2 + label: + en_US: gpt2 + label: + en_US: Tokenizer + human_description: + en_US: | + · cl100k_base --- gpt-4, gpt-3.5-turbo, gpt-3.5 + · o200k_base --- gpt-4o, gpt-4o-mini + · p50k_base --- text-davinci-003, text-davinci-002 + · r50k_base --- text-davinci-001, text-curie-001 + · p50k_edit --- text-davinci-edit-001, code-davinci-edit-001 + · gpt2 --- gpt-2 + form: form diff --git a/api/core/tools/provider/builtin/json_process/_assets/icon.svg b/api/core/tools/provider/builtin/json_process/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..b123983836962acd9d8289cbd879d221c77dc6a8 --- /dev/null +++ b/api/core/tools/provider/builtin/json_process/_assets/icon.svg @@ -0,0 +1,358 @@ + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/json_process/json_process.py b/api/core/tools/provider/builtin/json_process/json_process.py new file mode 100644 index 0000000000000000000000000000000000000000..10746210b5c6520b1903adaf857405773594a1e2 --- /dev/null +++ b/api/core/tools/provider/builtin/json_process/json_process.py @@ -0,0 +1,16 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.json_process.tools.parse import JSONParseTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class JsonExtractProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + JSONParseTool().invoke( + user_id="", + tool_parameters={"content": '{"name": "John", "age": 30, "city": "New York"}', "json_filter": "$.name"}, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/json_process/json_process.yaml b/api/core/tools/provider/builtin/json_process/json_process.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c7896bbea7a69f407212b4b839a6b98c51f607c8 --- /dev/null +++ b/api/core/tools/provider/builtin/json_process/json_process.yaml @@ -0,0 +1,14 @@ +identity: + author: Mingwei Zhang + name: json_process + label: + en_US: JSON Process + zh_Hans: JSON 处理 + pt_BR: JSON Process + description: + en_US: Tools for processing JSON content using jsonpath_ng + zh_Hans: 利用 jsonpath_ng 处理 JSON 内容的工具 + pt_BR: Tools for processing JSON content using jsonpath_ng + icon: icon.svg + tags: + - utilities diff --git a/api/core/tools/provider/builtin/json_process/tools/delete.py b/api/core/tools/provider/builtin/json_process/tools/delete.py new file mode 100644 index 0000000000000000000000000000000000000000..06f6cacd5d6126b713f8523c610ccffecc9e03b5 --- /dev/null +++ b/api/core/tools/provider/builtin/json_process/tools/delete.py @@ -0,0 +1,61 @@ +import json +from typing import Any, Union + +from jsonpath_ng import parse # type: ignore + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class JSONDeleteTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Invoke the JSON delete tool + """ + # Get content + content = tool_parameters.get("content", "") + if not content: + return self.create_text_message("Invalid parameter content") + + # Get query + query = tool_parameters.get("query", "") + if not query: + return self.create_text_message("Invalid parameter query") + + ensure_ascii = tool_parameters.get("ensure_ascii", True) + try: + result = self._delete(content, query, ensure_ascii) + return self.create_text_message(str(result)) + except Exception as e: + return self.create_text_message(f"Failed to delete JSON content: {str(e)}") + + def _delete(self, origin_json: str, query: str, ensure_ascii: bool) -> str: + try: + input_data = json.loads(origin_json) + expr = parse("$." + query.lstrip("$.")) # Ensure query path starts with $ + + matches = expr.find(input_data) + + if not matches: + return json.dumps(input_data, ensure_ascii=ensure_ascii) # No changes if no matches found + + for match in matches: + if isinstance(match.context.value, dict): + # Delete key from dictionary + del match.context.value[match.path.fields[-1]] + elif isinstance(match.context.value, list): + # Remove item from list + match.context.value.remove(match.value) + else: + # For other cases, we might want to set to None or remove the parent key + parent = match.context.parent + if parent: + del parent.value[match.path.fields[-1]] + + return json.dumps(input_data, ensure_ascii=ensure_ascii) + except Exception as e: + raise Exception(f"Delete operation failed: {str(e)}") diff --git a/api/core/tools/provider/builtin/json_process/tools/delete.yaml b/api/core/tools/provider/builtin/json_process/tools/delete.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4d390e40d172326d29e90d036f1cbd37bf6dde6f --- /dev/null +++ b/api/core/tools/provider/builtin/json_process/tools/delete.yaml @@ -0,0 +1,52 @@ +identity: + name: json_delete + author: Mingwei Zhang + label: + en_US: JSON Delete + zh_Hans: JSON 删除 + pt_BR: JSON Delete +description: + human: + en_US: A tool for deleting JSON content + zh_Hans: 一个删除 JSON 内容的工具 + pt_BR: A tool for deleting JSON content + llm: A tool for deleting JSON content +parameters: + - name: content + type: string + required: true + label: + en_US: JSON content + zh_Hans: JSON 内容 + pt_BR: JSON content + human_description: + en_US: JSON content to be processed + zh_Hans: 待处理的 JSON 内容 + pt_BR: JSON content to be processed + llm_description: JSON content to be processed + form: llm + - name: query + type: string + required: true + label: + en_US: Query + zh_Hans: 查询 + pt_BR: Query + human_description: + en_US: JSONPath query to locate the element to delete + zh_Hans: 用于定位要删除元素的 JSONPath 查询 + pt_BR: JSONPath query to locate the element to delete + llm_description: JSONPath query to locate the element to delete + form: llm + - name: ensure_ascii + type: boolean + default: true + label: + en_US: Ensure ASCII + zh_Hans: 确保 ASCII + pt_BR: Ensure ASCII + human_description: + en_US: Ensure the JSON output is ASCII encoded + zh_Hans: 确保输出的 JSON 是 ASCII 编码 + pt_BR: Ensure the JSON output is ASCII encoded + form: form diff --git a/api/core/tools/provider/builtin/json_process/tools/insert.py b/api/core/tools/provider/builtin/json_process/tools/insert.py new file mode 100644 index 0000000000000000000000000000000000000000..e825329a6d8f6191e740cbe99621fe989fc11b4a --- /dev/null +++ b/api/core/tools/provider/builtin/json_process/tools/insert.py @@ -0,0 +1,105 @@ +import json +from typing import Any, Union + +from jsonpath_ng import parse # type: ignore + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class JSONParseTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + # get content + content = tool_parameters.get("content", "") + if not content: + return self.create_text_message("Invalid parameter content") + + # get query + query = tool_parameters.get("query", "") + if not query: + return self.create_text_message("Invalid parameter query") + + # get new value + new_value = tool_parameters.get("new_value", "") + if not new_value: + return self.create_text_message("Invalid parameter new_value") + + # get insert position + index = tool_parameters.get("index") + + # get create path + create_path = tool_parameters.get("create_path", False) + + # get value decode. + # if true, it will be decoded to an dict + value_decode = tool_parameters.get("value_decode", False) + + ensure_ascii = tool_parameters.get("ensure_ascii", True) + try: + result = self._insert(content, query, new_value, ensure_ascii, value_decode, index, create_path) + return self.create_text_message(str(result)) + except Exception: + return self.create_text_message("Failed to insert JSON content") + + def _insert( + self, origin_json, query, new_value, ensure_ascii: bool, value_decode: bool, index=None, create_path=False + ): + try: + input_data = json.loads(origin_json) + expr = parse(query) + if value_decode is True: + try: + new_value = json.loads(new_value) + except json.JSONDecodeError: + return "Cannot decode new value to json object" + + matches = expr.find(input_data) + + if not matches and create_path: + # create new path + path_parts = query.strip("$").strip(".").split(".") + current = input_data + for i, part in enumerate(path_parts): + if "[" in part and "]" in part: + # process array index + array_name, index = part.split("[") + index = int(index.rstrip("]")) + if array_name not in current: + current[array_name] = [] + while len(current[array_name]) <= index: + current[array_name].append({}) + current = current[array_name][index] + else: + if i == len(path_parts) - 1: + current[part] = new_value + elif part not in current: + current[part] = {} + current = current[part] + else: + for match in matches: + if isinstance(match.value, dict): + # insert new value into dict + if isinstance(new_value, dict): + match.value.update(new_value) + else: + raise ValueError("Cannot insert non-dict value into dict") + elif isinstance(match.value, list): + # insert new value into list + if index is None: + match.value.append(new_value) + else: + match.value.insert(int(index), new_value) + else: + # replace old value with new value + match.full_path.update(input_data, new_value) + + return json.dumps(input_data, ensure_ascii=ensure_ascii) + except Exception as e: + return str(e) diff --git a/api/core/tools/provider/builtin/json_process/tools/insert.yaml b/api/core/tools/provider/builtin/json_process/tools/insert.yaml new file mode 100644 index 0000000000000000000000000000000000000000..21b51312dab6b3f4b8728f70ed227c0fab263eac --- /dev/null +++ b/api/core/tools/provider/builtin/json_process/tools/insert.yaml @@ -0,0 +1,101 @@ +identity: + name: json_insert + author: Mingwei Zhang + label: + en_US: JSON Insert + zh_Hans: JSON 插入 + pt_BR: JSON Insert +description: + human: + en_US: A tool for inserting JSON content + zh_Hans: 一个插入 JSON 内容的工具 + pt_BR: A tool for inserting JSON content + llm: A tool for inserting JSON content +parameters: + - name: content + type: string + required: true + label: + en_US: JSON content + zh_Hans: JSON 内容 + pt_BR: JSON content + human_description: + en_US: JSON content + zh_Hans: JSON 内容 + pt_BR: JSON content + llm_description: JSON content to be processed + form: llm + - name: query + type: string + required: true + label: + en_US: Query + zh_Hans: 查询 + pt_BR: Query + human_description: + en_US: Object to insert + zh_Hans: 待插入的对象 + pt_BR: Object to insert + llm_description: JSONPath query to locate the element to insert + form: llm + - name: new_value + type: string + required: true + label: + en_US: New Value + zh_Hans: 新值 + pt_BR: New Value + human_description: + en_US: New Value + zh_Hans: 插入的新值 + pt_BR: New Value + llm_description: New Value to insert + form: llm + - name: value_decode + type: boolean + default: false + label: + en_US: Decode Value + zh_Hans: 解码值 + pt_BR: Decode Value + human_description: + en_US: Whether to decode the value to a JSON object + zh_Hans: 是否将值解码为 JSON 对象 + pt_BR: Whether to decode the value to a JSON object + form: form + - name: create_path + type: select + required: true + default: "False" + label: + en_US: Whether to create a path + zh_Hans: 是否创建路径 + pt_BR: Whether to create a path + human_description: + en_US: Whether to create a path when the path does not exist + zh_Hans: 查询路径不存在时是否创建路径 + pt_BR: Whether to create a path when the path does not exist + options: + - value: "True" + label: + en_US: "Yes" + zh_Hans: 是 + pt_BR: "Yes" + - value: "False" + label: + en_US: "No" + zh_Hans: 否 + pt_BR: "No" + form: form + - name: ensure_ascii + type: boolean + default: true + label: + en_US: Ensure ASCII + zh_Hans: 确保 ASCII + pt_BR: Ensure ASCII + human_description: + en_US: Ensure the JSON output is ASCII encoded + zh_Hans: 确保输出的 JSON 是 ASCII 编码 + pt_BR: Ensure the JSON output is ASCII encoded + form: form diff --git a/api/core/tools/provider/builtin/json_process/tools/parse.py b/api/core/tools/provider/builtin/json_process/tools/parse.py new file mode 100644 index 0000000000000000000000000000000000000000..193017ba9a7c539d23f17429bb271b1c5a10dcae --- /dev/null +++ b/api/core/tools/provider/builtin/json_process/tools/parse.py @@ -0,0 +1,56 @@ +import json +from typing import Any, Union + +from jsonpath_ng import parse # type: ignore + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class JSONParseTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + # get content + content = tool_parameters.get("content", "") + if not content: + return self.create_text_message("Invalid parameter content") + + # get json filter + json_filter = tool_parameters.get("json_filter", "") + if not json_filter: + return self.create_text_message("Invalid parameter json_filter") + + ensure_ascii = tool_parameters.get("ensure_ascii", True) + try: + result = self._extract(content, json_filter, ensure_ascii) + return self.create_text_message(str(result)) + except Exception: + return self.create_text_message("Failed to extract JSON content") + + # Extract data from JSON content + def _extract(self, content: str, json_filter: str, ensure_ascii: bool) -> str: + try: + input_data = json.loads(content) + expr = parse(json_filter) + result = [match.value for match in expr.find(input_data)] + + if not result: + return "" + + if len(result) == 1: + result = result[0] + + if isinstance(result, dict | list): + return json.dumps(result, ensure_ascii=ensure_ascii) + elif isinstance(result, str | int | float | bool) or result is None: + return str(result) + else: + return repr(result) + except Exception as e: + return str(e) diff --git a/api/core/tools/provider/builtin/json_process/tools/parse.yaml b/api/core/tools/provider/builtin/json_process/tools/parse.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c35f4eac0775adbb2b55d5e79a92cd4676e23850 --- /dev/null +++ b/api/core/tools/provider/builtin/json_process/tools/parse.yaml @@ -0,0 +1,52 @@ +identity: + name: parse + author: Mingwei Zhang + label: + en_US: JSON Parse + zh_Hans: JSON 解析 + pt_BR: JSON Parse +description: + human: + en_US: A tool for extracting JSON objects + zh_Hans: 一个解析JSON对象的工具 + pt_BR: A tool for extracting JSON objects + llm: A tool for extracting JSON objects +parameters: + - name: content + type: string + required: true + label: + en_US: JSON data + zh_Hans: JSON数据 + pt_BR: JSON data + human_description: + en_US: JSON data + zh_Hans: JSON数据 + pt_BR: JSON数据 + llm_description: JSON data to be processed + form: llm + - name: json_filter + type: string + required: true + label: + en_US: JSON filter + zh_Hans: JSON解析对象 + pt_BR: JSON filter + human_description: + en_US: JSON fields to be parsed + zh_Hans: 需要解析的 JSON 字段 + pt_BR: JSON fields to be parsed + llm_description: JSON fields to be parsed + form: llm + - name: ensure_ascii + type: boolean + default: true + label: + en_US: Ensure ASCII + zh_Hans: 确保 ASCII + pt_BR: Ensure ASCII + human_description: + en_US: Ensure the JSON output is ASCII encoded + zh_Hans: 确保输出的 JSON 是 ASCII 编码 + pt_BR: Ensure the JSON output is ASCII encoded + form: form diff --git a/api/core/tools/provider/builtin/json_process/tools/replace.py b/api/core/tools/provider/builtin/json_process/tools/replace.py new file mode 100644 index 0000000000000000000000000000000000000000..feca0d8a7c278382f6861b11cfa35c9c04e76ab5 --- /dev/null +++ b/api/core/tools/provider/builtin/json_process/tools/replace.py @@ -0,0 +1,129 @@ +import json +from typing import Any, Union + +from jsonpath_ng import parse # type: ignore + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class JSONReplaceTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + # get content + content = tool_parameters.get("content", "") + if not content: + return self.create_text_message("Invalid parameter content") + + # get query + query = tool_parameters.get("query", "") + if not query: + return self.create_text_message("Invalid parameter query") + + # get replace value + replace_value = tool_parameters.get("replace_value", "") + if not replace_value: + return self.create_text_message("Invalid parameter replace_value") + + # get replace model + replace_model = tool_parameters.get("replace_model", "") + if not replace_model: + return self.create_text_message("Invalid parameter replace_model") + + # get value decode. + # if true, it will be decoded to an dict + value_decode = tool_parameters.get("value_decode", False) + + ensure_ascii = tool_parameters.get("ensure_ascii", True) + try: + if replace_model == "pattern": + # get replace pattern + replace_pattern = tool_parameters.get("replace_pattern", "") + if not replace_pattern: + return self.create_text_message("Invalid parameter replace_pattern") + result = self._replace_pattern( + content, query, replace_pattern, replace_value, ensure_ascii, value_decode + ) + elif replace_model == "key": + result = self._replace_key(content, query, replace_value, ensure_ascii) + elif replace_model == "value": + result = self._replace_value(content, query, replace_value, ensure_ascii, value_decode) + return self.create_text_message(str(result)) + except Exception: + return self.create_text_message("Failed to replace JSON content") + + # Replace pattern + def _replace_pattern( + self, content: str, query: str, replace_pattern: str, replace_value: str, ensure_ascii: bool, value_decode: bool + ) -> str: + try: + input_data = json.loads(content) + expr = parse(query) + + matches = expr.find(input_data) + + for match in matches: + new_value = match.value.replace(replace_pattern, replace_value) + if value_decode is True: + try: + new_value = json.loads(new_value) + except json.JSONDecodeError: + return "Cannot decode replace value to json object" + + match.full_path.update(input_data, new_value) + + return json.dumps(input_data, ensure_ascii=ensure_ascii) + except Exception as e: + return str(e) + + # Replace key + def _replace_key(self, content: str, query: str, replace_value: str, ensure_ascii: bool) -> str: + try: + input_data = json.loads(content) + expr = parse(query) + + matches = expr.find(input_data) + + for match in matches: + parent = match.context.value + if isinstance(parent, dict): + old_key = match.path.fields[0] + if old_key in parent: + value = parent.pop(old_key) + parent[replace_value] = value + elif isinstance(parent, list): + for item in parent: + if isinstance(item, dict) and old_key in item: + value = item.pop(old_key) + item[replace_value] = value + return json.dumps(input_data, ensure_ascii=ensure_ascii) + except Exception as e: + return str(e) + + # Replace value + def _replace_value( + self, content: str, query: str, replace_value: str, ensure_ascii: bool, value_decode: bool + ) -> str: + try: + input_data = json.loads(content) + expr = parse(query) + if value_decode is True: + try: + replace_value = json.loads(replace_value) + except json.JSONDecodeError: + return "Cannot decode replace value to json object" + + matches = expr.find(input_data) + + for match in matches: + match.full_path.update(input_data, replace_value) + + return json.dumps(input_data, ensure_ascii=ensure_ascii) + except Exception as e: + return str(e) diff --git a/api/core/tools/provider/builtin/json_process/tools/replace.yaml b/api/core/tools/provider/builtin/json_process/tools/replace.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ae238b1fbcd05e5f69bc4042cddf6415c719fa8b --- /dev/null +++ b/api/core/tools/provider/builtin/json_process/tools/replace.yaml @@ -0,0 +1,119 @@ +identity: + name: json_replace + author: Mingwei Zhang + label: + en_US: JSON Replace + zh_Hans: JSON 替换 + pt_BR: JSON Replace +description: + human: + en_US: A tool for replacing JSON content + zh_Hans: 一个替换 JSON 内容的工具 + pt_BR: A tool for replacing JSON content + llm: A tool for replacing JSON content +parameters: + - name: content + type: string + required: true + label: + en_US: JSON content + zh_Hans: JSON 内容 + pt_BR: JSON content + human_description: + en_US: JSON content + zh_Hans: JSON 内容 + pt_BR: JSON content + llm_description: JSON content to be processed + form: llm + - name: query + type: string + required: true + label: + en_US: Query + zh_Hans: 查询 + pt_BR: Query + human_description: + en_US: Query + zh_Hans: 查询 + pt_BR: Query + llm_description: JSONPath query to locate the element to replace + form: llm + - name: replace_pattern + type: string + required: false + label: + en_US: String to be replaced + zh_Hans: 待替换字符串 + pt_BR: String to be replaced + human_description: + en_US: String to be replaced + zh_Hans: 待替换字符串 + pt_BR: String to be replaced + llm_description: String to be replaced + form: llm + - name: replace_value + type: string + required: true + label: + en_US: Replace Value + zh_Hans: 替换值 + pt_BR: Replace Value + human_description: + en_US: New Value + zh_Hans: 新值 + pt_BR: New Value + llm_description: New Value to replace + form: llm + - name: value_decode + type: boolean + default: false + label: + en_US: Decode Value + zh_Hans: 解码值 + pt_BR: Decode Value + human_description: + en_US: Whether to decode the value to a JSON object (Does not apply to replace key) + zh_Hans: 是否将值解码为 JSON 对象 (不适用于键替换) + pt_BR: Whether to decode the value to a JSON object (Does not apply to replace key) + form: form + - name: replace_model + type: select + required: true + default: pattern + label: + en_US: Replace Model + zh_Hans: 替换模式 + pt_BR: Replace Model + human_description: + en_US: Replace Model + zh_Hans: 替换模式 + pt_BR: Replace Model + options: + - value: key + label: + en_US: replace key + zh_Hans: 键替换 + pt_BR: replace key + - value: value + label: + en_US: replace value + zh_Hans: 值替换 + pt_BR: replace value + - value: pattern + label: + en_US: replace string + zh_Hans: 字符串替换 + pt_BR: replace string + form: form + - name: ensure_ascii + type: boolean + default: true + label: + en_US: Ensure ASCII + zh_Hans: 确保 ASCII + pt_BR: Ensure ASCII + human_description: + en_US: Ensure the JSON output is ASCII encoded + zh_Hans: 确保输出的 JSON 是 ASCII 编码 + pt_BR: Ensure the JSON output is ASCII encoded + form: form diff --git a/api/core/tools/provider/builtin/judge0ce/_assets/icon.svg b/api/core/tools/provider/builtin/judge0ce/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..3e7e33da6e8b25b7fb55c6049a5009bd0936e0cf --- /dev/null +++ b/api/core/tools/provider/builtin/judge0ce/_assets/icon.svg @@ -0,0 +1,21 @@ + + + + + + + + + + diff --git a/api/core/tools/provider/builtin/judge0ce/judge0ce.py b/api/core/tools/provider/builtin/judge0ce/judge0ce.py new file mode 100644 index 0000000000000000000000000000000000000000..50db74dd9ebced8efd59e869745d7fb6344affe4 --- /dev/null +++ b/api/core/tools/provider/builtin/judge0ce/judge0ce.py @@ -0,0 +1,23 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.judge0ce.tools.executeCode import ExecuteCodeTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class Judge0CEProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + ExecuteCodeTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "source_code": "print('hello world')", + "language_id": 71, + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/judge0ce/judge0ce.yaml b/api/core/tools/provider/builtin/judge0ce/judge0ce.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9ff8aaac6debc6dc80055301d4d935d3fdb579bf --- /dev/null +++ b/api/core/tools/provider/builtin/judge0ce/judge0ce.yaml @@ -0,0 +1,32 @@ +identity: + author: Richards Tu + name: judge0ce + label: + en_US: Judge0 CE + zh_Hans: Judge0 CE + pt_BR: Judge0 CE + description: + en_US: Judge0 CE is an open-source code execution system. Support various languages, including C, C++, Java, Python, Ruby, etc. + zh_Hans: Judge0 CE 是一个开源的代码执行系统。支持多种语言,包括 C、C++、Java、Python、Ruby 等。 + pt_BR: Judge0 CE é um sistema de execução de código de código aberto. Suporta várias linguagens, incluindo C, C++, Java, Python, Ruby, etc. + icon: icon.svg + tags: + - utilities + - other +credentials_for_provider: + X-RapidAPI-Key: + type: secret-input + required: true + label: + en_US: RapidAPI Key + zh_Hans: RapidAPI Key + pt_BR: RapidAPI Key + help: + en_US: RapidAPI Key is required to access the Judge0 CE API. + zh_Hans: RapidAPI Key 是访问 Judge0 CE API 所必需的。 + pt_BR: RapidAPI Key é necessário para acessar a API do Judge0 CE. + placeholder: + en_US: Enter your RapidAPI Key + zh_Hans: 输入你的 RapidAPI Key + pt_BR: Insira sua RapidAPI Key + url: https://rapidapi.com/judge0-official/api/judge0-ce diff --git a/api/core/tools/provider/builtin/judge0ce/tools/executeCode.py b/api/core/tools/provider/builtin/judge0ce/tools/executeCode.py new file mode 100644 index 0000000000000000000000000000000000000000..b8d654ff639575a9db95798a94df0f8861ecc30a --- /dev/null +++ b/api/core/tools/provider/builtin/judge0ce/tools/executeCode.py @@ -0,0 +1,61 @@ +import json +from typing import Any, Union + +import requests +from httpx import post + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class ExecuteCodeTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + api_key = self.runtime.credentials["X-RapidAPI-Key"] + + url = "https://judge0-ce.p.rapidapi.com/submissions" + + querystring = {"base64_encoded": "false", "fields": "*"} + + headers = { + "Content-Type": "application/json", + "X-RapidAPI-Key": api_key, + "X-RapidAPI-Host": "judge0-ce.p.rapidapi.com", + } + + payload = { + "language_id": tool_parameters["language_id"], + "source_code": tool_parameters["source_code"], + "stdin": tool_parameters.get("stdin", ""), + "expected_output": tool_parameters.get("expected_output", ""), + "additional_files": tool_parameters.get("additional_files", ""), + } + + response = post(url, data=json.dumps(payload), headers=headers, params=querystring) + + if response.status_code != 201: + raise Exception(response.text) + + token = response.json()["token"] + + url = f"https://judge0-ce.p.rapidapi.com/submissions/{token}" + headers = {"X-RapidAPI-Key": api_key} + + response = requests.get(url, headers=headers) + if response.status_code == 200: + result = response.json() + return self.create_text_message( + text=f"stdout: {result.get('stdout', '')}\n" + f"stderr: {result.get('stderr', '')}\n" + f"compile_output: {result.get('compile_output', '')}\n" + f"message: {result.get('message', '')}\n" + f"status: {result['status']['description']}\n" + f"time: {result.get('time', '')} seconds\n" + f"memory: {result.get('memory', '')} bytes" + ) + else: + return self.create_text_message(text=f"Error retrieving submission details: {response.text}") diff --git a/api/core/tools/provider/builtin/judge0ce/tools/executeCode.yaml b/api/core/tools/provider/builtin/judge0ce/tools/executeCode.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a8c0776f40185e02202a851d0eee5292046f1ea3 --- /dev/null +++ b/api/core/tools/provider/builtin/judge0ce/tools/executeCode.yaml @@ -0,0 +1,67 @@ +identity: + name: submitCodeExecutionTask + author: Richards Tu + label: + en_US: Submit Code Execution Task to Judge0 CE and get execution result. + zh_Hans: 提交代码执行任务到 Judge0 CE 并获取执行结果。 +description: + human: + en_US: A tool for executing code and getting the result. + zh_Hans: 一个用于执行代码并获取结果的工具。 + llm: This tool is used for executing code and getting the result. +parameters: + - name: source_code + type: string + required: true + label: + en_US: Source Code + zh_Hans: 源代码 + human_description: + en_US: The source code to be executed. + zh_Hans: 要执行的源代码。 + llm_description: The source code to be executed. + form: llm + - name: language_id + type: number + required: true + label: + en_US: Language ID + zh_Hans: 语言 ID + human_description: + en_US: The ID of the language in which the source code is written. + zh_Hans: 源代码所使用的语言的 ID。 + llm_description: The ID of the language in which the source code is written. For example, 50 for C++, 71 for Python, etc. + form: llm + - name: stdin + type: string + required: false + label: + en_US: Standard Input + zh_Hans: 标准输入 + human_description: + en_US: The standard input to be provided to the program. + zh_Hans: 提供给程序的标准输入。 + llm_description: The standard input to be provided to the program. Optional. + form: llm + - name: expected_output + type: string + required: false + label: + en_US: Expected Output + zh_Hans: 期望输出 + human_description: + en_US: The expected output of the program. Used for comparison in some scenarios. + zh_Hans: 程序的期望输出。在某些场景下用于比较。 + llm_description: The expected output of the program. Used for comparison in some scenarios. Optional. + form: llm + - name: additional_files + type: string + required: false + label: + en_US: Additional Files + zh_Hans: 附加文件 + human_description: + en_US: Base64 encoded additional files for the submission. + zh_Hans: 提交的 Base64 编码的附加文件。 + llm_description: Base64 encoded additional files for the submission. Optional. + form: llm diff --git a/api/core/tools/provider/builtin/lark_base/_assets/icon.png b/api/core/tools/provider/builtin/lark_base/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..036e586772ef502791c73db50681267a69bab98d Binary files /dev/null and b/api/core/tools/provider/builtin/lark_base/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/lark_base/lark_base.py b/api/core/tools/provider/builtin/lark_base/lark_base.py new file mode 100644 index 0000000000000000000000000000000000000000..de9b36831198445c41c8fbd78b1d7306f3a99022 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/lark_base.py @@ -0,0 +1,7 @@ +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController +from core.tools.utils.lark_api_utils import lark_auth + + +class LarkBaseProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + lark_auth(credentials) diff --git a/api/core/tools/provider/builtin/lark_base/lark_base.yaml b/api/core/tools/provider/builtin/lark_base/lark_base.yaml new file mode 100644 index 0000000000000000000000000000000000000000..200b2e22cfa558d96f670631a6f8171e3fcf315a --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/lark_base.yaml @@ -0,0 +1,36 @@ +identity: + author: Doug Lea + name: lark_base + label: + en_US: Lark Base + zh_Hans: Lark 多维表格 + description: + en_US: | + Lark base, requires the following permissions: bitable:app. + zh_Hans: | + Lark 多维表格,需要开通以下权限: bitable:app。 + icon: icon.png + tags: + - social + - productivity +credentials_for_provider: + app_id: + type: text-input + required: true + label: + en_US: APP ID + placeholder: + en_US: Please input your Lark app id + zh_Hans: 请输入你的 Lark app id + help: + en_US: Get your app_id and app_secret from Lark + zh_Hans: 从 Lark 获取您的 app_id 和 app_secret + url: https://open.larksuite.com/app + app_secret: + type: secret-input + required: true + label: + en_US: APP Secret + placeholder: + en_US: Please input your app secret + zh_Hans: 请输入你的 Lark app secret diff --git a/api/core/tools/provider/builtin/lark_base/tools/add_records.py b/api/core/tools/provider/builtin/lark_base/tools/add_records.py new file mode 100644 index 0000000000000000000000000000000000000000..c46898062a8cc2950773a427d681149165c96eec --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/tools/add_records.py @@ -0,0 +1,21 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class AddRecordsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + app_token = tool_parameters.get("app_token") + table_id = tool_parameters.get("table_id") + table_name = tool_parameters.get("table_name") + records = tool_parameters.get("records") + user_id_type = tool_parameters.get("user_id_type", "open_id") + + res = client.add_records(app_token, table_id, table_name, records, user_id_type) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_base/tools/add_records.yaml b/api/core/tools/provider/builtin/lark_base/tools/add_records.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f2a93490dc0c3103b9fe4de29f54d4bdd6db5bdc --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/tools/add_records.yaml @@ -0,0 +1,91 @@ +identity: + name: add_records + author: Doug Lea + label: + en_US: Add Records + zh_Hans: 新增多条记录 +description: + human: + en_US: Add Multiple Records to Multidimensional Table + zh_Hans: 在多维表格数据表中新增多条记录 + llm: A tool for adding multiple records to a multidimensional table. (在多维表格数据表中新增多条记录) +parameters: + - name: app_token + type: string + required: true + label: + en_US: app_token + zh_Hans: app_token + human_description: + en_US: Unique identifier for the multidimensional table, supports inputting document URL. + zh_Hans: 多维表格的唯一标识符,支持输入文档 URL。 + llm_description: 多维表格的唯一标识符,支持输入文档 URL。 + form: llm + + - name: table_id + type: string + required: false + label: + en_US: table_id + zh_Hans: table_id + human_description: + en_US: Unique identifier for the multidimensional table data, either table_id or table_name must be provided, cannot be empty simultaneously. + zh_Hans: 多维表格数据表的唯一标识符,table_id 和 table_name 至少需要提供一个,不能同时为空。 + llm_description: 多维表格数据表的唯一标识符,table_id 和 table_name 至少需要提供一个,不能同时为空。 + form: llm + + - name: table_name + type: string + required: false + label: + en_US: table_name + zh_Hans: table_name + human_description: + en_US: Name of the multidimensional table data, either table_name or table_id must be provided, cannot be empty simultaneously. + zh_Hans: 多维表格数据表的名称,table_name 和 table_id 至少需要提供一个,不能同时为空。 + llm_description: 多维表格数据表的名称,table_name 和 table_id 至少需要提供一个,不能同时为空。 + form: llm + + - name: records + type: string + required: true + label: + en_US: records + zh_Hans: 记录列表 + human_description: + en_US: | + List of records to be added in this request. Example value: [{"multi-line-text":"text content","single_select":"option 1","date":1674206443000}] + For supported field types, refer to the integration guide (https://open.larkoffice.com/document/server-docs/docs/bitable-v1/notification). For data structures of different field types, refer to the data structure overview (https://open.larkoffice.com/document/server-docs/docs/bitable-v1/bitable-structure). + zh_Hans: | + 本次请求将要新增的记录列表,示例值:[{"多行文本":"文本内容","单选":"选项 1","日期":1674206443000}]。 + 当前接口支持的字段类型请参考接入指南(https://open.larkoffice.com/document/server-docs/docs/bitable-v1/notification),不同类型字段的数据结构请参考数据结构概述(https://open.larkoffice.com/document/server-docs/docs/bitable-v1/bitable-structure)。 + llm_description: | + 本次请求将要新增的记录列表,示例值:[{"多行文本":"文本内容","单选":"选项 1","日期":1674206443000}]。 + 当前接口支持的字段类型请参考接入指南(https://open.larkoffice.com/document/server-docs/docs/bitable-v1/notification),不同类型字段的数据结构请参考数据结构概述(https://open.larkoffice.com/document/server-docs/docs/bitable-v1/bitable-structure)。 + form: llm + + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form diff --git a/api/core/tools/provider/builtin/lark_base/tools/create_base.py b/api/core/tools/provider/builtin/lark_base/tools/create_base.py new file mode 100644 index 0000000000000000000000000000000000000000..a857c6ced6f94bcf32479fdc718ce2f952678c8f --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/tools/create_base.py @@ -0,0 +1,18 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class CreateBaseTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + name = tool_parameters.get("name") + folder_token = tool_parameters.get("folder_token") + + res = client.create_base(name, folder_token) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_base/tools/create_base.yaml b/api/core/tools/provider/builtin/lark_base/tools/create_base.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e622edf3362ba42e64e0053a1beaa687fd215045 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/tools/create_base.yaml @@ -0,0 +1,42 @@ +identity: + name: create_base + author: Doug Lea + label: + en_US: Create Base + zh_Hans: 创建多维表格 +description: + human: + en_US: Create Multidimensional Table in Specified Directory + zh_Hans: 在指定目录下创建多维表格 + llm: A tool for creating a multidimensional table in a specified directory. (在指定目录下创建多维表格) +parameters: + - name: name + type: string + required: false + label: + en_US: name + zh_Hans: 多维表格 App 名字 + human_description: + en_US: | + Name of the multidimensional table App. Example value: "A new multidimensional table". + zh_Hans: 多维表格 App 名字,示例值:"一篇新的多维表格"。 + llm_description: 多维表格 App 名字,示例值:"一篇新的多维表格"。 + form: llm + + - name: folder_token + type: string + required: false + label: + en_US: folder_token + zh_Hans: 多维表格 App 归属文件夹 + human_description: + en_US: | + Folder where the multidimensional table App belongs. Default is empty, meaning the table will be created in the root directory of the cloud space. Example values: Lf8uf6BoAlWkUfdGtpMjUV0PpZd or https://lark-japan.jp.larksuite.com/drive/folder/Lf8uf6BoAlWkUfdGtpMjUV0PpZd. + The folder_token must be an existing folder and supports inputting folder token or folder URL. + zh_Hans: | + 多维表格 App 归属文件夹。默认为空,表示多维表格将被创建在云空间根目录。示例值: Lf8uf6BoAlWkUfdGtpMjUV0PpZd 或者 https://lark-japan.jp.larksuite.com/drive/folder/Lf8uf6BoAlWkUfdGtpMjUV0PpZd。 + folder_token 必须是已存在的文件夹,支持输入文件夹 token 或者文件夹 URL。 + llm_description: | + 多维表格 App 归属文件夹。默认为空,表示多维表格将被创建在云空间根目录。示例值: Lf8uf6BoAlWkUfdGtpMjUV0PpZd 或者 https://lark-japan.jp.larksuite.com/drive/folder/Lf8uf6BoAlWkUfdGtpMjUV0PpZd。 + folder_token 必须是已存在的文件夹,支持输入文件夹 token 或者文件夹 URL。 + form: llm diff --git a/api/core/tools/provider/builtin/lark_base/tools/create_table.py b/api/core/tools/provider/builtin/lark_base/tools/create_table.py new file mode 100644 index 0000000000000000000000000000000000000000..aff7e715b73a733a925f813086ef961cbc8950a1 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/tools/create_table.py @@ -0,0 +1,20 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class CreateTableTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + app_token = tool_parameters.get("app_token") + table_name = tool_parameters.get("table_name") + default_view_name = tool_parameters.get("default_view_name") + fields = tool_parameters.get("fields") + + res = client.create_table(app_token, table_name, default_view_name, fields) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_base/tools/create_table.yaml b/api/core/tools/provider/builtin/lark_base/tools/create_table.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8b1007b9a531663b87846dfcdd3f075b81929420 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/tools/create_table.yaml @@ -0,0 +1,61 @@ +identity: + name: create_table + author: Doug Lea + label: + en_US: Create Table + zh_Hans: 新增数据表 +description: + human: + en_US: Add a Data Table to Multidimensional Table + zh_Hans: 在多维表格中新增一个数据表 + llm: A tool for adding a data table to a multidimensional table. (在多维表格中新增一个数据表) +parameters: + - name: app_token + type: string + required: true + label: + en_US: app_token + zh_Hans: app_token + human_description: + en_US: Unique identifier for the multidimensional table, supports inputting document URL. + zh_Hans: 多维表格的唯一标识符,支持输入文档 URL。 + llm_description: 多维表格的唯一标识符,支持输入文档 URL。 + form: llm + + - name: table_name + type: string + required: true + label: + en_US: Table Name + zh_Hans: 数据表名称 + human_description: + en_US: | + The name of the data table, length range: 1 character to 100 characters. + zh_Hans: 数据表名称,长度范围:1 字符 ~ 100 字符。 + llm_description: 数据表名称,长度范围:1 字符 ~ 100 字符。 + form: llm + + - name: default_view_name + type: string + required: false + label: + en_US: Default View Name + zh_Hans: 默认表格视图的名称 + human_description: + en_US: The name of the default table view, defaults to "Table" if not filled. + zh_Hans: 默认表格视图的名称,不填则默认为"表格"。 + llm_description: 默认表格视图的名称,不填则默认为"表格"。 + form: llm + + - name: fields + type: string + required: true + label: + en_US: Initial Fields + zh_Hans: 初始字段 + human_description: + en_US: | + Initial fields of the data table, format: [ { "field_name": "Multi-line Text","type": 1 },{ "field_name": "Number","type": 2 },{ "field_name": "Single Select","type": 3 },{ "field_name": "Multiple Select","type": 4 },{ "field_name": "Date","type": 5 } ]. For field details, refer to: https://open.larkoffice.com/document/server-docs/docs/bitable-v1/app-table-field/guide + zh_Hans: 数据表的初始字段,格式为:[{"field_name":"多行文本","type":1},{"field_name":"数字","type":2},{"field_name":"单选","type":3},{"field_name":"多选","type":4},{"field_name":"日期","type":5}]。字段详情参考:https://open.larkoffice.com/document/server-docs/docs/bitable-v1/app-table-field/guide + llm_description: 数据表的初始字段,格式为:[{"field_name":"多行文本","type":1},{"field_name":"数字","type":2},{"field_name":"单选","type":3},{"field_name":"多选","type":4},{"field_name":"日期","type":5}]。字段详情参考:https://open.larkoffice.com/document/server-docs/docs/bitable-v1/app-table-field/guide + form: llm diff --git a/api/core/tools/provider/builtin/lark_base/tools/delete_records.py b/api/core/tools/provider/builtin/lark_base/tools/delete_records.py new file mode 100644 index 0000000000000000000000000000000000000000..1b0a7470505e4d76664b4c3c5c0ca74249ab7ee2 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/tools/delete_records.py @@ -0,0 +1,20 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class DeleteRecordsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + app_token = tool_parameters.get("app_token") + table_id = tool_parameters.get("table_id") + table_name = tool_parameters.get("table_name") + record_ids = tool_parameters.get("record_ids") + + res = client.delete_records(app_token, table_id, table_name, record_ids) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_base/tools/delete_records.yaml b/api/core/tools/provider/builtin/lark_base/tools/delete_records.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c30ebd630ce9d835a78fa77724cecf16acfe5dbe --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/tools/delete_records.yaml @@ -0,0 +1,86 @@ +identity: + name: delete_records + author: Doug Lea + label: + en_US: Delete Records + zh_Hans: 删除多条记录 +description: + human: + en_US: Delete Multiple Records from Multidimensional Table + zh_Hans: 删除多维表格数据表中的多条记录 + llm: A tool for deleting multiple records from a multidimensional table. (删除多维表格数据表中的多条记录) +parameters: + - name: app_token + type: string + required: true + label: + en_US: app_token + zh_Hans: app_token + human_description: + en_US: Unique identifier for the multidimensional table, supports inputting document URL. + zh_Hans: 多维表格的唯一标识符,支持输入文档 URL。 + llm_description: 多维表格的唯一标识符,支持输入文档 URL。 + form: llm + + - name: table_id + type: string + required: false + label: + en_US: table_id + zh_Hans: table_id + human_description: + en_US: Unique identifier for the multidimensional table data, either table_id or table_name must be provided, cannot be empty simultaneously. + zh_Hans: 多维表格数据表的唯一标识符,table_id 和 table_name 至少需要提供一个,不能同时为空。 + llm_description: 多维表格数据表的唯一标识符,table_id 和 table_name 至少需要提供一个,不能同时为空。 + form: llm + + - name: table_name + type: string + required: false + label: + en_US: table_name + zh_Hans: table_name + human_description: + en_US: Name of the multidimensional table data, either table_name or table_id must be provided, cannot be empty simultaneously. + zh_Hans: 多维表格数据表的名称,table_name 和 table_id 至少需要提供一个,不能同时为空。 + llm_description: 多维表格数据表的名称,table_name 和 table_id 至少需要提供一个,不能同时为空。 + form: llm + + - name: record_ids + type: string + required: true + label: + en_US: Record IDs + zh_Hans: 记录 ID 列表 + human_description: + en_US: | + List of IDs for the records to be deleted, example value: ["recwNXzPQv"]. + zh_Hans: 删除的多条记录 ID 列表,示例值:["recwNXzPQv"]。 + llm_description: 删除的多条记录 ID 列表,示例值:["recwNXzPQv"]。 + form: llm + + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form diff --git a/api/core/tools/provider/builtin/lark_base/tools/delete_tables.py b/api/core/tools/provider/builtin/lark_base/tools/delete_tables.py new file mode 100644 index 0000000000000000000000000000000000000000..e0ecae2f1750505c6eaaec819e0c3a88f55a8602 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/tools/delete_tables.py @@ -0,0 +1,19 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class DeleteTablesTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + app_token = tool_parameters.get("app_token") + table_ids = tool_parameters.get("table_ids") + table_names = tool_parameters.get("table_names") + + res = client.delete_tables(app_token, table_ids, table_names) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_base/tools/delete_tables.yaml b/api/core/tools/provider/builtin/lark_base/tools/delete_tables.yaml new file mode 100644 index 0000000000000000000000000000000000000000..498126eae53302d088f275d9f3fc71c9b6cff378 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/tools/delete_tables.yaml @@ -0,0 +1,49 @@ +identity: + name: delete_tables + author: Doug Lea + label: + en_US: Delete Tables + zh_Hans: 删除数据表 +description: + human: + en_US: Batch Delete Data Tables from Multidimensional Table + zh_Hans: 批量删除多维表格中的数据表 + llm: A tool for batch deleting data tables from a multidimensional table. (批量删除多维表格中的数据表) +parameters: + - name: app_token + type: string + required: true + label: + en_US: app_token + zh_Hans: app_token + human_description: + en_US: Unique identifier for the multidimensional table, supports inputting document URL. + zh_Hans: 多维表格的唯一标识符,支持输入文档 URL。 + llm_description: 多维表格的唯一标识符,支持输入文档 URL。 + form: llm + + - name: table_ids + type: string + required: false + label: + en_US: Table IDs + zh_Hans: 数据表 ID + human_description: + en_US: | + IDs of the tables to be deleted. Each operation supports deleting up to 50 tables. Example: ["tbl1TkhyTWDkSoZ3"]. Ensure that either table_ids or table_names is not empty. + zh_Hans: 待删除的数据表的 ID,每次操作最多支持删除 50 个数据表。示例值:["tbl1TkhyTWDkSoZ3"]。请确保 table_ids 和 table_names 至少有一个不为空。 + llm_description: 待删除的数据表的 ID,每次操作最多支持删除 50 个数据表。示例值:["tbl1TkhyTWDkSoZ3"]。请确保 table_ids 和 table_names 至少有一个不为空。 + form: llm + + - name: table_names + type: string + required: false + label: + en_US: Table Names + zh_Hans: 数据表名称 + human_description: + en_US: | + Names of the tables to be deleted. Each operation supports deleting up to 50 tables. Example: ["Table1", "Table2"]. Ensure that either table_names or table_ids is not empty. + zh_Hans: 待删除的数据表的名称,每次操作最多支持删除 50 个数据表。示例值:["数据表1", "数据表2"]。请确保 table_names 和 table_ids 至少有一个不为空。 + llm_description: 待删除的数据表的名称,每次操作最多支持删除 50 个数据表。示例值:["数据表1", "数据表2"]。请确保 table_names 和 table_ids 至少有一个不为空。 + form: llm diff --git a/api/core/tools/provider/builtin/lark_base/tools/get_base_info.py b/api/core/tools/provider/builtin/lark_base/tools/get_base_info.py new file mode 100644 index 0000000000000000000000000000000000000000..2c23248b88765a55311695c89359f9d85b004070 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/tools/get_base_info.py @@ -0,0 +1,17 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class GetBaseInfoTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + app_token = tool_parameters.get("app_token") + + res = client.get_base_info(app_token) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_base/tools/get_base_info.yaml b/api/core/tools/provider/builtin/lark_base/tools/get_base_info.yaml new file mode 100644 index 0000000000000000000000000000000000000000..eb0e7a26c06a557b6335b82b7f46825cfabf8b5f --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/tools/get_base_info.yaml @@ -0,0 +1,23 @@ +identity: + name: get_base_info + author: Doug Lea + label: + en_US: Get Base Info + zh_Hans: 获取多维表格元数据 +description: + human: + en_US: Get Metadata Information of Specified Multidimensional Table + zh_Hans: 获取指定多维表格的元数据信息 + llm: A tool for getting metadata information of a specified multidimensional table. (获取指定多维表格的元数据信息) +parameters: + - name: app_token + type: string + required: true + label: + en_US: app_token + zh_Hans: app_token + human_description: + en_US: Unique identifier for the multidimensional table, supports inputting document URL. + zh_Hans: 多维表格的唯一标识符,支持输入文档 URL。 + llm_description: 多维表格的唯一标识符,支持输入文档 URL。 + form: llm diff --git a/api/core/tools/provider/builtin/lark_base/tools/list_tables.py b/api/core/tools/provider/builtin/lark_base/tools/list_tables.py new file mode 100644 index 0000000000000000000000000000000000000000..55b706854b27356cc2e5224efb272782ecc19c15 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/tools/list_tables.py @@ -0,0 +1,19 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class ListTablesTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + app_token = tool_parameters.get("app_token") + page_token = tool_parameters.get("page_token") + page_size = tool_parameters.get("page_size", 20) + + res = client.list_tables(app_token, page_token, page_size) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_base/tools/list_tables.yaml b/api/core/tools/provider/builtin/lark_base/tools/list_tables.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7571519039bd242132cb3655378be85e60461111 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/tools/list_tables.yaml @@ -0,0 +1,50 @@ +identity: + name: list_tables + author: Doug Lea + label: + en_US: List Tables + zh_Hans: 列出数据表 +description: + human: + en_US: Get All Data Tables under Multidimensional Table + zh_Hans: 获取多维表格下的所有数据表 + llm: A tool for getting all data tables under a multidimensional table. (获取多维表格下的所有数据表) +parameters: + - name: app_token + type: string + required: true + label: + en_US: app_token + zh_Hans: app_token + human_description: + en_US: Unique identifier for the multidimensional table, supports inputting document URL. + zh_Hans: 多维表格的唯一标识符,支持输入文档 URL。 + llm_description: 多维表格的唯一标识符,支持输入文档 URL。 + form: llm + + - name: page_size + type: number + required: false + default: 20 + label: + en_US: page_size + zh_Hans: 分页大小 + human_description: + en_US: | + Page size, default value: 20, maximum value: 100. + zh_Hans: 分页大小,默认值:20,最大值:100。 + llm_description: 分页大小,默认值:20,最大值:100。 + form: form + + - name: page_token + type: string + required: false + label: + en_US: page_token + zh_Hans: 分页标记 + human_description: + en_US: | + Page token, leave empty for the first request to start from the beginning; a new page_token will be returned if there are more items in the paginated query results, which can be used for the next traversal. Example value: "tblsRc9GRRXKqhvW". + zh_Hans: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。示例值:"tblsRc9GRRXKqhvW"。 + llm_description: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。示例值:"tblsRc9GRRXKqhvW"。 + form: llm diff --git a/api/core/tools/provider/builtin/lark_base/tools/read_records.py b/api/core/tools/provider/builtin/lark_base/tools/read_records.py new file mode 100644 index 0000000000000000000000000000000000000000..5cf25aad848dfac54adfb6236bddf3c63d90532a --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/tools/read_records.py @@ -0,0 +1,21 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class ReadRecordsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + app_token = tool_parameters.get("app_token") + table_id = tool_parameters.get("table_id") + table_name = tool_parameters.get("table_name") + record_ids = tool_parameters.get("record_ids") + user_id_type = tool_parameters.get("user_id_type", "open_id") + + res = client.read_records(app_token, table_id, table_name, record_ids, user_id_type) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_base/tools/read_records.yaml b/api/core/tools/provider/builtin/lark_base/tools/read_records.yaml new file mode 100644 index 0000000000000000000000000000000000000000..911e667cfc90adf5890378f50c858376f58b569d --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/tools/read_records.yaml @@ -0,0 +1,86 @@ +identity: + name: read_records + author: Doug Lea + label: + en_US: Read Records + zh_Hans: 批量获取记录 +description: + human: + en_US: Batch Retrieve Records from Multidimensional Table + zh_Hans: 批量获取多维表格数据表中的记录信息 + llm: A tool for batch retrieving records from a multidimensional table, supporting up to 100 records per call. (批量获取多维表格数据表中的记录信息,单次调用最多支持查询 100 条记录) + +parameters: + - name: app_token + type: string + required: true + label: + en_US: app_token + zh_Hans: app_token + human_description: + en_US: Unique identifier for the multidimensional table, supports inputting document URL. + zh_Hans: 多维表格的唯一标识符,支持输入文档 URL。 + llm_description: 多维表格的唯一标识符,支持输入文档 URL。 + form: llm + + - name: table_id + type: string + required: false + label: + en_US: table_id + zh_Hans: table_id + human_description: + en_US: Unique identifier for the multidimensional table data, either table_id or table_name must be provided, cannot be empty simultaneously. + zh_Hans: 多维表格数据表的唯一标识符,table_id 和 table_name 至少需要提供一个,不能同时为空。 + llm_description: 多维表格数据表的唯一标识符,table_id 和 table_name 至少需要提供一个,不能同时为空。 + form: llm + + - name: table_name + type: string + required: false + label: + en_US: table_name + zh_Hans: table_name + human_description: + en_US: Name of the multidimensional table data, either table_name or table_id must be provided, cannot be empty simultaneously. + zh_Hans: 多维表格数据表的名称,table_name 和 table_id 至少需要提供一个,不能同时为空。 + llm_description: 多维表格数据表的名称,table_name 和 table_id 至少需要提供一个,不能同时为空。 + form: llm + + - name: record_ids + type: string + required: true + label: + en_US: record_ids + zh_Hans: 记录 ID 列表 + human_description: + en_US: List of record IDs, which can be obtained by calling the "Query Records API". + zh_Hans: 记录 ID 列表,可以通过调用"查询记录接口"获取。 + llm_description: 记录 ID 列表,可以通过调用"查询记录接口"获取。 + form: llm + + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form diff --git a/api/core/tools/provider/builtin/lark_base/tools/search_records.py b/api/core/tools/provider/builtin/lark_base/tools/search_records.py new file mode 100644 index 0000000000000000000000000000000000000000..9b0abcf067951af274911b8534d524885846b723 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/tools/search_records.py @@ -0,0 +1,39 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class SearchRecordsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + app_token = tool_parameters.get("app_token") + table_id = tool_parameters.get("table_id") + table_name = tool_parameters.get("table_name") + view_id = tool_parameters.get("view_id") + field_names = tool_parameters.get("field_names") + sort = tool_parameters.get("sort") + filters = tool_parameters.get("filter") + page_token = tool_parameters.get("page_token") + automatic_fields = tool_parameters.get("automatic_fields", False) + user_id_type = tool_parameters.get("user_id_type", "open_id") + page_size = tool_parameters.get("page_size", 20) + + res = client.search_record( + app_token, + table_id, + table_name, + view_id, + field_names, + sort, + filters, + page_token, + automatic_fields, + user_id_type, + page_size, + ) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_base/tools/search_records.yaml b/api/core/tools/provider/builtin/lark_base/tools/search_records.yaml new file mode 100644 index 0000000000000000000000000000000000000000..edd86ab9d69686788db359d12d8fad6d77fed268 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/tools/search_records.yaml @@ -0,0 +1,163 @@ +identity: + name: search_records + author: Doug Lea + label: + en_US: Search Records + zh_Hans: 查询记录 +description: + human: + en_US: Query records in a multidimensional table, up to 500 rows per query. + zh_Hans: 查询多维表格数据表中的记录,单次最多查询 500 行记录。 + llm: A tool for querying records in a multidimensional table, up to 500 rows per query. (查询多维表格数据表中的记录,单次最多查询 500 行记录) +parameters: + - name: app_token + type: string + required: true + label: + en_US: app_token + zh_Hans: app_token + human_description: + en_US: Unique identifier for the multidimensional table, supports inputting document URL. + zh_Hans: 多维表格的唯一标识符,支持输入文档 URL。 + llm_description: 多维表格的唯一标识符,支持输入文档 URL。 + form: llm + + - name: table_id + type: string + required: false + label: + en_US: table_id + zh_Hans: table_id + human_description: + en_US: Unique identifier for the multidimensional table data, either table_id or table_name must be provided, cannot be empty simultaneously. + zh_Hans: 多维表格数据表的唯一标识符,table_id 和 table_name 至少需要提供一个,不能同时为空。 + llm_description: 多维表格数据表的唯一标识符,table_id 和 table_name 至少需要提供一个,不能同时为空。 + form: llm + + - name: table_name + type: string + required: false + label: + en_US: table_name + zh_Hans: table_name + human_description: + en_US: Name of the multidimensional table data, either table_name or table_id must be provided, cannot be empty simultaneously. + zh_Hans: 多维表格数据表的名称,table_name 和 table_id 至少需要提供一个,不能同时为空。 + llm_description: 多维表格数据表的名称,table_name 和 table_id 至少需要提供一个,不能同时为空。 + form: llm + + - name: view_id + type: string + required: false + label: + en_US: view_id + zh_Hans: 视图唯一标识 + human_description: + en_US: | + Unique identifier for a view in a multidimensional table. It can be found in the URL's query parameter with the key 'view'. For example: https://lark-japan.jp.larksuite.com/base/XXX0bfYEraW5OWsbhcFjEqj6pxh?table=tbl5I6jqwz8wBRMv&view=vewW5zXVEU. + zh_Hans: 多维表格中视图的唯一标识,可在多维表格的 URL 地址栏中找到,query 参数中 key 为 view 的部分。例如:https://lark-japan.jp.larksuite.com/base/XXX0bfYEraW5OWsbhcFjEqj6pxh?table=tbl5I6jqwz8wBRMv&view=vewW5zXVEU。 + llm_description: 多维表格中视图的唯一标识,可在多维表格的 URL 地址栏中找到,query 参数中 key 为 view 的部分。例如:https://lark-japan.jp.larksuite.com/base/XXX0bfYEraW5OWsbhcFjEqj6pxh?table=tbl5I6jqwz8wBRMv&view=vewW5zXVEU。 + form: llm + + - name: field_names + type: string + required: false + label: + en_US: field_names + zh_Hans: 字段名称 + human_description: + en_US: | + Field names to specify which fields to include in the returned records. Example value: ["Field1", "Field2"]. + zh_Hans: 字段名称,用于指定本次查询返回记录中包含的字段。示例值:["字段1","字段2"]。 + llm_description: 字段名称,用于指定本次查询返回记录中包含的字段。示例值:["字段1","字段2"]。 + form: llm + + - name: sort + type: string + required: false + label: + en_US: sort + zh_Hans: 排序条件 + human_description: + en_US: | + Sorting conditions, for example: [{"field_name":"Multiline Text","desc":true}]. + zh_Hans: 排序条件,例如:[{"field_name":"多行文本","desc":true}]。 + llm_description: 排序条件,例如:[{"field_name":"多行文本","desc":true}]。 + form: llm + + - name: filter + type: string + required: false + label: + en_US: filter + zh_Hans: 筛选条件 + human_description: + en_US: Object containing filter information. For details on how to fill in the filter, refer to the record filter parameter guide (https://open.larkoffice.com/document/uAjLw4CM/ukTMukTMukTM/reference/bitable-v1/app-table-record/record-filter-guide). + zh_Hans: 包含条件筛选信息的对象。了解如何填写 filter,参考记录筛选参数填写指南(https://open.larkoffice.com/document/uAjLw4CM/ukTMukTMukTM/reference/bitable-v1/app-table-record/record-filter-guide)。 + llm_description: 包含条件筛选信息的对象。了解如何填写 filter,参考记录筛选参数填写指南(https://open.larkoffice.com/document/uAjLw4CM/ukTMukTMukTM/reference/bitable-v1/app-table-record/record-filter-guide)。 + form: llm + + - name: automatic_fields + type: boolean + required: false + label: + en_US: automatic_fields + zh_Hans: automatic_fields + human_description: + en_US: Whether to return automatically calculated fields. Default is false, meaning they are not returned. + zh_Hans: 是否返回自动计算的字段。默认为 false,表示不返回。 + llm_description: 是否返回自动计算的字段。默认为 false,表示不返回。 + form: form + + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form + + - name: page_size + type: number + required: false + default: 20 + label: + en_US: page_size + zh_Hans: 分页大小 + human_description: + en_US: | + Page size, default value: 20, maximum value: 500. + zh_Hans: 分页大小,默认值:20,最大值:500。 + llm_description: 分页大小,默认值:20,最大值:500。 + form: form + + - name: page_token + type: string + required: false + label: + en_US: page_token + zh_Hans: 分页标记 + human_description: + en_US: | + Page token, leave empty for the first request to start from the beginning; a new page_token will be returned if there are more items in the paginated query results, which can be used for the next traversal. Example value: "tblsRc9GRRXKqhvW". + zh_Hans: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。示例值:"tblsRc9GRRXKqhvW"。 + llm_description: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。示例值:"tblsRc9GRRXKqhvW"。 + form: llm diff --git a/api/core/tools/provider/builtin/lark_base/tools/update_records.py b/api/core/tools/provider/builtin/lark_base/tools/update_records.py new file mode 100644 index 0000000000000000000000000000000000000000..7c263df2bb031c21668098369b4eaa8e3b489872 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/tools/update_records.py @@ -0,0 +1,21 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class UpdateRecordsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + app_token = tool_parameters.get("app_token") + table_id = tool_parameters.get("table_id") + table_name = tool_parameters.get("table_name") + records = tool_parameters.get("records") + user_id_type = tool_parameters.get("user_id_type", "open_id") + + res = client.update_records(app_token, table_id, table_name, records, user_id_type) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_base/tools/update_records.yaml b/api/core/tools/provider/builtin/lark_base/tools/update_records.yaml new file mode 100644 index 0000000000000000000000000000000000000000..68117e7136789225bf75724bbdf59c9fbccfbd36 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_base/tools/update_records.yaml @@ -0,0 +1,91 @@ +identity: + name: update_records + author: Doug Lea + label: + en_US: Update Records + zh_Hans: 更新多条记录 +description: + human: + en_US: Update Multiple Records in Multidimensional Table + zh_Hans: 更新多维表格数据表中的多条记录 + llm: A tool for updating multiple records in a multidimensional table. (更新多维表格数据表中的多条记录) +parameters: + - name: app_token + type: string + required: true + label: + en_US: app_token + zh_Hans: app_token + human_description: + en_US: Unique identifier for the multidimensional table, supports inputting document URL. + zh_Hans: 多维表格的唯一标识符,支持输入文档 URL。 + llm_description: 多维表格的唯一标识符,支持输入文档 URL。 + form: llm + + - name: table_id + type: string + required: false + label: + en_US: table_id + zh_Hans: table_id + human_description: + en_US: Unique identifier for the multidimensional table data, either table_id or table_name must be provided, cannot be empty simultaneously. + zh_Hans: 多维表格数据表的唯一标识符,table_id 和 table_name 至少需要提供一个,不能同时为空。 + llm_description: 多维表格数据表的唯一标识符,table_id 和 table_name 至少需要提供一个,不能同时为空。 + form: llm + + - name: table_name + type: string + required: false + label: + en_US: table_name + zh_Hans: table_name + human_description: + en_US: Name of the multidimensional table data, either table_name or table_id must be provided, cannot be empty simultaneously. + zh_Hans: 多维表格数据表的名称,table_name 和 table_id 至少需要提供一个,不能同时为空。 + llm_description: 多维表格数据表的名称,table_name 和 table_id 至少需要提供一个,不能同时为空。 + form: llm + + - name: records + type: string + required: true + label: + en_US: records + zh_Hans: 记录列表 + human_description: + en_US: | + List of records to be updated in this request. Example value: [{"fields":{"multi-line-text":"text content","single_select":"option 1","date":1674206443000},"record_id":"recupK4f4RM5RX"}]. + For supported field types, refer to the integration guide (https://open.larkoffice.com/document/server-docs/docs/bitable-v1/notification). For data structures of different field types, refer to the data structure overview (https://open.larkoffice.com/document/server-docs/docs/bitable-v1/bitable-structure). + zh_Hans: | + 本次请求将要更新的记录列表,示例值:[{"fields":{"多行文本":"文本内容","单选":"选项 1","日期":1674206443000},"record_id":"recupK4f4RM5RX"}]。 + 当前接口支持的字段类型请参考接入指南(https://open.larkoffice.com/document/server-docs/docs/bitable-v1/notification),不同类型字段的数据结构请参考数据结构概述(https://open.larkoffice.com/document/server-docs/docs/bitable-v1/bitable-structure)。 + llm_description: | + 本次请求将要更新的记录列表,示例值:[{"fields":{"多行文本":"文本内容","单选":"选项 1","日期":1674206443000},"record_id":"recupK4f4RM5RX"}]。 + 当前接口支持的字段类型请参考接入指南(https://open.larkoffice.com/document/server-docs/docs/bitable-v1/notification),不同类型字段的数据结构请参考数据结构概述(https://open.larkoffice.com/document/server-docs/docs/bitable-v1/bitable-structure)。 + form: llm + + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form diff --git a/api/core/tools/provider/builtin/lark_calendar/_assets/icon.png b/api/core/tools/provider/builtin/lark_calendar/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..2a934747a98c6680065941bcd31d2400da1eaf23 Binary files /dev/null and b/api/core/tools/provider/builtin/lark_calendar/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/lark_calendar/lark_calendar.py b/api/core/tools/provider/builtin/lark_calendar/lark_calendar.py new file mode 100644 index 0000000000000000000000000000000000000000..871de69cc15b3986a719dd9f30e5962b01148f57 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_calendar/lark_calendar.py @@ -0,0 +1,7 @@ +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController +from core.tools.utils.lark_api_utils import lark_auth + + +class LarkCalendarProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + lark_auth(credentials) diff --git a/api/core/tools/provider/builtin/lark_calendar/lark_calendar.yaml b/api/core/tools/provider/builtin/lark_calendar/lark_calendar.yaml new file mode 100644 index 0000000000000000000000000000000000000000..72c41e36c0ebd363a9a5fcb8f17a2bef3a4925b9 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_calendar/lark_calendar.yaml @@ -0,0 +1,36 @@ +identity: + author: Doug Lea + name: lark_calendar + label: + en_US: Lark Calendar + zh_Hans: Lark 日历 + description: + en_US: | + Lark calendar, requires the following permissions: calendar:calendar:read、calendar:calendar、contact:user.id:readonly. + zh_Hans: | + Lark 日历,需要开通以下权限: calendar:calendar:read、calendar:calendar、contact:user.id:readonly。 + icon: icon.png + tags: + - social + - productivity +credentials_for_provider: + app_id: + type: text-input + required: true + label: + en_US: APP ID + placeholder: + en_US: Please input your Lark app id + zh_Hans: 请输入你的 Lark app id + help: + en_US: Get your app_id and app_secret from Lark + zh_Hans: 从 Lark 获取您的 app_id 和 app_secret + url: https://open.larksuite.com/app + app_secret: + type: secret-input + required: true + label: + en_US: APP Secret + placeholder: + en_US: Please input your app secret + zh_Hans: 请输入你的 Lark app secret diff --git a/api/core/tools/provider/builtin/lark_calendar/tools/add_event_attendees.py b/api/core/tools/provider/builtin/lark_calendar/tools/add_event_attendees.py new file mode 100644 index 0000000000000000000000000000000000000000..f5929893ddfe245d0936347c30d6afef73e6b123 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_calendar/tools/add_event_attendees.py @@ -0,0 +1,20 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class AddEventAttendeesTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + event_id = tool_parameters.get("event_id") + attendee_phone_or_email = tool_parameters.get("attendee_phone_or_email") + need_notification = tool_parameters.get("need_notification", True) + + res = client.add_event_attendees(event_id, attendee_phone_or_email, need_notification) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_calendar/tools/add_event_attendees.yaml b/api/core/tools/provider/builtin/lark_calendar/tools/add_event_attendees.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9d7a1319072d6f4c351612680e53a8d150d26dbd --- /dev/null +++ b/api/core/tools/provider/builtin/lark_calendar/tools/add_event_attendees.yaml @@ -0,0 +1,54 @@ +identity: + name: add_event_attendees + author: Doug Lea + label: + en_US: Add Event Attendees + zh_Hans: 添加日程参会人 +description: + human: + en_US: Add Event Attendees + zh_Hans: 添加日程参会人 + llm: A tool for adding attendees to events in Lark. (在 Lark 中添加日程参会人) +parameters: + - name: event_id + type: string + required: true + label: + en_US: Event ID + zh_Hans: 日程 ID + human_description: + en_US: | + The ID of the event, which will be returned when the event is created. For example: fb2a6406-26d6-4c8d-a487-6f0246c94d2f_0. + zh_Hans: | + 创建日程时会返回日程 ID。例如: fb2a6406-26d6-4c8d-a487-6f0246c94d2f_0。 + llm_description: | + 日程 ID,创建日程时会返回日程 ID。例如: fb2a6406-26d6-4c8d-a487-6f0246c94d2f_0。 + form: llm + + - name: need_notification + type: boolean + required: false + default: true + label: + en_US: Need Notification + zh_Hans: 是否需要通知 + human_description: + en_US: | + Whether to send a Bot notification to attendees. true: send, false: do not send. + zh_Hans: | + 是否给参与人发送 Bot 通知,true: 发送,false: 不发送。 + llm_description: | + 是否给参与人发送 Bot 通知,true: 发送,false: 不发送。 + form: form + + - name: attendee_phone_or_email + type: string + required: true + label: + en_US: Attendee Phone or Email + zh_Hans: 参会人电话或邮箱 + human_description: + en_US: The list of attendee emails or phone numbers, separated by commas. + zh_Hans: 日程参会人邮箱或者手机号列表,使用逗号分隔。 + llm_description: 日程参会人邮箱或者手机号列表,使用逗号分隔。 + form: llm diff --git a/api/core/tools/provider/builtin/lark_calendar/tools/create_event.py b/api/core/tools/provider/builtin/lark_calendar/tools/create_event.py new file mode 100644 index 0000000000000000000000000000000000000000..8a0726008c3f8b839d504ffd3695ba11c2aff469 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_calendar/tools/create_event.py @@ -0,0 +1,26 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class CreateEventTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + summary = tool_parameters.get("summary") + description = tool_parameters.get("description") + start_time = tool_parameters.get("start_time") + end_time = tool_parameters.get("end_time") + attendee_ability = tool_parameters.get("attendee_ability") + need_notification = tool_parameters.get("need_notification", True) + auto_record = tool_parameters.get("auto_record", False) + + res = client.create_event( + summary, description, start_time, end_time, attendee_ability, need_notification, auto_record + ) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_calendar/tools/create_event.yaml b/api/core/tools/provider/builtin/lark_calendar/tools/create_event.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b738736e630fa51496264aee58ac868f82c4d44a --- /dev/null +++ b/api/core/tools/provider/builtin/lark_calendar/tools/create_event.yaml @@ -0,0 +1,119 @@ +identity: + name: create_event + author: Doug Lea + label: + en_US: Create Event + zh_Hans: 创建日程 +description: + human: + en_US: Create Event + zh_Hans: 创建日程 + llm: A tool for creating events in Lark.(创建 Lark 日程) +parameters: + - name: summary + type: string + required: false + label: + en_US: Summary + zh_Hans: 日程标题 + human_description: + en_US: The title of the event. If not filled, the event title will display (No Subject). + zh_Hans: 日程标题,若不填则日程标题显示 (无主题)。 + llm_description: 日程标题,若不填则日程标题显示 (无主题)。 + form: llm + + - name: description + type: string + required: false + label: + en_US: Description + zh_Hans: 日程描述 + human_description: + en_US: The description of the event. + zh_Hans: 日程描述。 + llm_description: 日程描述。 + form: llm + + - name: need_notification + type: boolean + required: false + default: true + label: + en_US: Need Notification + zh_Hans: 是否发送通知 + human_description: + en_US: | + Whether to send a bot message when the event is created, true: send, false: do not send. + zh_Hans: 创建日程时是否发送 bot 消息,true:发送,false:不发送。 + llm_description: 创建日程时是否发送 bot 消息,true:发送,false:不发送。 + form: form + + - name: start_time + type: string + required: true + label: + en_US: Start Time + zh_Hans: 开始时间 + human_description: + en_US: | + The start time of the event, format: 2006-01-02 15:04:05. + zh_Hans: 日程开始时间,格式:2006-01-02 15:04:05。 + llm_description: 日程开始时间,格式:2006-01-02 15:04:05。 + form: llm + + - name: end_time + type: string + required: true + label: + en_US: End Time + zh_Hans: 结束时间 + human_description: + en_US: | + The end time of the event, format: 2006-01-02 15:04:05. + zh_Hans: 日程结束时间,格式:2006-01-02 15:04:05。 + llm_description: 日程结束时间,格式:2006-01-02 15:04:05。 + form: llm + + - name: attendee_ability + type: select + required: false + options: + - value: none + label: + en_US: none + zh_Hans: 无 + - value: can_see_others + label: + en_US: can_see_others + zh_Hans: 可以查看参与人列表 + - value: can_invite_others + label: + en_US: can_invite_others + zh_Hans: 可以邀请其它参与人 + - value: can_modify_event + label: + en_US: can_modify_event + zh_Hans: 可以编辑日程 + default: "none" + label: + en_US: attendee_ability + zh_Hans: 参会人权限 + human_description: + en_US: Attendee ability, optional values are none, can_see_others, can_invite_others, can_modify_event, with a default value of none. + zh_Hans: 参会人权限,可选值有无、可以查看参与人列表、可以邀请其它参与人、可以编辑日程,默认值为无。 + llm_description: 参会人权限,可选值有无、可以查看参与人列表、可以邀请其它参与人、可以编辑日程,默认值为无。 + form: form + + - name: auto_record + type: boolean + required: false + default: false + label: + en_US: Auto Record + zh_Hans: 自动录制 + human_description: + en_US: | + Whether to enable automatic recording, true: enabled, automatically record when the meeting starts; false: not enabled. + zh_Hans: 是否开启自动录制,true:开启,会议开始后自动录制;false:不开启。 + llm_description: 是否开启自动录制,true:开启,会议开始后自动录制;false:不开启。 + form: form diff --git a/api/core/tools/provider/builtin/lark_calendar/tools/delete_event.py b/api/core/tools/provider/builtin/lark_calendar/tools/delete_event.py new file mode 100644 index 0000000000000000000000000000000000000000..0e4ceac5e5d0708bef5c0a805a191ffff70ec1ea --- /dev/null +++ b/api/core/tools/provider/builtin/lark_calendar/tools/delete_event.py @@ -0,0 +1,19 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class DeleteEventTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + event_id = tool_parameters.get("event_id") + need_notification = tool_parameters.get("need_notification", True) + + res = client.delete_event(event_id, need_notification) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_calendar/tools/delete_event.yaml b/api/core/tools/provider/builtin/lark_calendar/tools/delete_event.yaml new file mode 100644 index 0000000000000000000000000000000000000000..cdd6d7e1bb024a03efc75f349e2b6746ce899425 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_calendar/tools/delete_event.yaml @@ -0,0 +1,38 @@ +identity: + name: delete_event + author: Doug Lea + label: + en_US: Delete Event + zh_Hans: 删除日程 +description: + human: + en_US: Delete Event + zh_Hans: 删除日程 + llm: A tool for deleting events in Lark.(在 Lark 中删除日程) +parameters: + - name: event_id + type: string + required: true + label: + en_US: Event ID + zh_Hans: 日程 ID + human_description: + en_US: | + The ID of the event, for example: e8b9791c-39ae-4908-8ad8-66b13159b9fb_0. + zh_Hans: 日程 ID,例如:e8b9791c-39ae-4908-8ad8-66b13159b9fb_0。 + llm_description: 日程 ID,例如:e8b9791c-39ae-4908-8ad8-66b13159b9fb_0。 + form: llm + + - name: need_notification + type: boolean + required: false + default: true + label: + en_US: Need Notification + zh_Hans: 是否需要通知 + human_description: + en_US: | + Indicates whether to send bot notifications to event participants upon deletion. true: send, false: do not send. + zh_Hans: 删除日程是否给日程参与人发送 bot 通知,true:发送,false:不发送。 + llm_description: 删除日程是否给日程参与人发送 bot 通知,true:发送,false:不发送。 + form: form diff --git a/api/core/tools/provider/builtin/lark_calendar/tools/get_primary_calendar.py b/api/core/tools/provider/builtin/lark_calendar/tools/get_primary_calendar.py new file mode 100644 index 0000000000000000000000000000000000000000..d315bf35f05d98bf969e10c656f24f25e1450840 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_calendar/tools/get_primary_calendar.py @@ -0,0 +1,18 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class GetPrimaryCalendarTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + user_id_type = tool_parameters.get("user_id_type", "open_id") + + res = client.get_primary_calendar(user_id_type) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_calendar/tools/get_primary_calendar.yaml b/api/core/tools/provider/builtin/lark_calendar/tools/get_primary_calendar.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fe615947700995c0b554b3791d1ab723dc0323bb --- /dev/null +++ b/api/core/tools/provider/builtin/lark_calendar/tools/get_primary_calendar.yaml @@ -0,0 +1,37 @@ +identity: + name: get_primary_calendar + author: Doug Lea + label: + en_US: Get Primary Calendar + zh_Hans: 查询主日历信息 +description: + human: + en_US: Get Primary Calendar + zh_Hans: 查询主日历信息 + llm: A tool for querying primary calendar information in Lark.(在 Lark 中查询主日历信息) +parameters: + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form diff --git a/api/core/tools/provider/builtin/lark_calendar/tools/list_events.py b/api/core/tools/provider/builtin/lark_calendar/tools/list_events.py new file mode 100644 index 0000000000000000000000000000000000000000..d74cc049d342309dab6ae26062ebc1509f17272a --- /dev/null +++ b/api/core/tools/provider/builtin/lark_calendar/tools/list_events.py @@ -0,0 +1,21 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class ListEventsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + start_time = tool_parameters.get("start_time") + end_time = tool_parameters.get("end_time") + page_token = tool_parameters.get("page_token") + page_size = tool_parameters.get("page_size") + + res = client.list_events(start_time, end_time, page_token, page_size) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_calendar/tools/list_events.yaml b/api/core/tools/provider/builtin/lark_calendar/tools/list_events.yaml new file mode 100644 index 0000000000000000000000000000000000000000..cef332f5272e5575665fc2984999798fa46c1b2e --- /dev/null +++ b/api/core/tools/provider/builtin/lark_calendar/tools/list_events.yaml @@ -0,0 +1,62 @@ +identity: + name: list_events + author: Doug Lea + label: + en_US: List Events + zh_Hans: 获取日程列表 +description: + human: + en_US: List Events + zh_Hans: 获取日程列表 + llm: A tool for listing events in Lark.(在 Lark 中获取日程列表) +parameters: + - name: start_time + type: string + required: false + label: + en_US: Start Time + zh_Hans: 开始时间 + human_description: + en_US: | + The start time, defaults to 0:00 of the current day if not provided, format: 2006-01-02 15:04:05. + zh_Hans: 开始时间,不传值时默认当天 0 点时间,格式为:2006-01-02 15:04:05。 + llm_description: 开始时间,不传值时默认当天 0 点时间,格式为:2006-01-02 15:04:05。 + form: llm + + - name: end_time + type: string + required: false + label: + en_US: End Time + zh_Hans: 结束时间 + human_description: + en_US: | + The end time, defaults to 23:59 of the current day if not provided, format: 2006-01-02 15:04:05. + zh_Hans: 结束时间,不传值时默认当天 23:59 分时间,格式为:2006-01-02 15:04:05。 + llm_description: 结束时间,不传值时默认当天 23:59 分时间,格式为:2006-01-02 15:04:05。 + form: llm + + - name: page_size + type: number + required: false + default: 50 + label: + en_US: Page Size + zh_Hans: 分页大小 + human_description: + en_US: The page size, i.e., the number of data entries returned in a single request. The default value is 50, and the value range is [50,1000]. + zh_Hans: 分页大小,即单次请求所返回的数据条目数。默认值为 50,取值范围为 [50,1000]。 + llm_description: 分页大小,即单次请求所返回的数据条目数。默认值为 50,取值范围为 [50,1000]。 + form: form + + - name: page_token + type: string + required: false + label: + en_US: Page Token + zh_Hans: 分页标记 + human_description: + en_US: The pagination token. Leave it blank for the first request, indicating to start traversing from the beginning; when the pagination query result has more items, a new page_token will be returned simultaneously, which can be used to obtain the query result in the next traversal. + zh_Hans: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + llm_description: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + form: llm diff --git a/api/core/tools/provider/builtin/lark_calendar/tools/search_events.py b/api/core/tools/provider/builtin/lark_calendar/tools/search_events.py new file mode 100644 index 0000000000000000000000000000000000000000..a20038e47dd430266f74fbd0d71c92dcd6820748 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_calendar/tools/search_events.py @@ -0,0 +1,23 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class SearchEventsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + query = tool_parameters.get("query") + start_time = tool_parameters.get("start_time") + end_time = tool_parameters.get("end_time") + page_token = tool_parameters.get("page_token") + user_id_type = tool_parameters.get("user_id_type", "open_id") + page_size = tool_parameters.get("page_size", 20) + + res = client.search_events(query, start_time, end_time, page_token, user_id_type, page_size) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_calendar/tools/search_events.yaml b/api/core/tools/provider/builtin/lark_calendar/tools/search_events.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4d4f8819c11e4d51338262b053bcf629e4b490bc --- /dev/null +++ b/api/core/tools/provider/builtin/lark_calendar/tools/search_events.yaml @@ -0,0 +1,100 @@ +identity: + name: search_events + author: Doug Lea + label: + en_US: Search Events + zh_Hans: 搜索日程 +description: + human: + en_US: Search Events + zh_Hans: 搜索日程 + llm: A tool for searching events in Lark.(在 Lark 中搜索日程) +parameters: + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form + + - name: query + type: string + required: true + label: + en_US: Query + zh_Hans: 搜索关键字 + human_description: + en_US: The search keyword used for fuzzy searching event names, with a maximum input of 200 characters. + zh_Hans: 用于模糊查询日程名称的搜索关键字,最大输入 200 字符。 + llm_description: 用于模糊查询日程名称的搜索关键字,最大输入 200 字符。 + form: llm + + - name: start_time + type: string + required: false + label: + en_US: Start Time + zh_Hans: 开始时间 + human_description: + en_US: | + The start time, defaults to 0:00 of the current day if not provided, format: 2006-01-02 15:04:05. + zh_Hans: 开始时间,不传值时默认当天 0 点时间,格式为:2006-01-02 15:04:05。 + llm_description: 开始时间,不传值时默认当天 0 点时间,格式为:2006-01-02 15:04:05。 + form: llm + + - name: end_time + type: string + required: false + label: + en_US: End Time + zh_Hans: 结束时间 + human_description: + en_US: | + The end time, defaults to 23:59 of the current day if not provided, format: 2006-01-02 15:04:05. + zh_Hans: 结束时间,不传值时默认当天 23:59 分时间,格式为:2006-01-02 15:04:05。 + llm_description: 结束时间,不传值时默认当天 23:59 分时间,格式为:2006-01-02 15:04:05。 + form: llm + + - name: page_size + type: number + required: false + default: 20 + label: + en_US: Page Size + zh_Hans: 分页大小 + human_description: + en_US: The page size, i.e., the number of data entries returned in a single request. The default value is 20, and the value range is [10,100]. + zh_Hans: 分页大小,即单次请求所返回的数据条目数。默认值为 20,取值范围为 [10,100]。 + llm_description: 分页大小,即单次请求所返回的数据条目数。默认值为 20,取值范围为 [10,100]。 + form: form + + - name: page_token + type: string + required: false + label: + en_US: Page Token + zh_Hans: 分页标记 + human_description: + en_US: The pagination token. Leave it blank for the first request, indicating to start traversing from the beginning; when the pagination query result has more items, a new page_token will be returned simultaneously, which can be used to obtain the query result in the next traversal. + zh_Hans: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + llm_description: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + form: llm diff --git a/api/core/tools/provider/builtin/lark_calendar/tools/update_event.py b/api/core/tools/provider/builtin/lark_calendar/tools/update_event.py new file mode 100644 index 0000000000000000000000000000000000000000..a04029377f6799e3816fb072b0cdaaeb7c1db96a --- /dev/null +++ b/api/core/tools/provider/builtin/lark_calendar/tools/update_event.py @@ -0,0 +1,24 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class UpdateEventTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + event_id = tool_parameters.get("event_id") + summary = tool_parameters.get("summary") + description = tool_parameters.get("description") + need_notification = tool_parameters.get("need_notification", True) + start_time = tool_parameters.get("start_time") + end_time = tool_parameters.get("end_time") + auto_record = tool_parameters.get("auto_record", False) + + res = client.update_event(event_id, summary, description, need_notification, start_time, end_time, auto_record) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_calendar/tools/update_event.yaml b/api/core/tools/provider/builtin/lark_calendar/tools/update_event.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b9992e5b03f944c66d5a954bc06f383ff1f7a4eb --- /dev/null +++ b/api/core/tools/provider/builtin/lark_calendar/tools/update_event.yaml @@ -0,0 +1,100 @@ +identity: + name: update_event + author: Doug Lea + label: + en_US: Update Event + zh_Hans: 更新日程 +description: + human: + en_US: Update Event + zh_Hans: 更新日程 + llm: A tool for updating events in Lark.(更新 Lark 中的日程) +parameters: + - name: event_id + type: string + required: true + label: + en_US: Event ID + zh_Hans: 日程 ID + human_description: + en_US: | + The ID of the event, for example: e8b9791c-39ae-4908-8ad8-66b13159b9fb_0. + zh_Hans: 日程 ID,例如:e8b9791c-39ae-4908-8ad8-66b13159b9fb_0。 + llm_description: 日程 ID,例如:e8b9791c-39ae-4908-8ad8-66b13159b9fb_0。 + form: llm + + - name: summary + type: string + required: false + label: + en_US: Summary + zh_Hans: 日程标题 + human_description: + en_US: The title of the event. + zh_Hans: 日程标题。 + llm_description: 日程标题。 + form: llm + + - name: description + type: string + required: false + label: + en_US: Description + zh_Hans: 日程描述 + human_description: + en_US: The description of the event. + zh_Hans: 日程描述。 + llm_description: 日程描述。 + form: llm + + - name: need_notification + type: boolean + required: false + label: + en_US: Need Notification + zh_Hans: 是否发送通知 + human_description: + en_US: | + Whether to send a bot message when the event is updated, true: send, false: do not send. + zh_Hans: 更新日程时是否发送 bot 消息,true:发送,false:不发送。 + llm_description: 更新日程时是否发送 bot 消息,true:发送,false:不发送。 + form: form + + - name: start_time + type: string + required: false + label: + en_US: Start Time + zh_Hans: 开始时间 + human_description: + en_US: | + The start time of the event, format: 2006-01-02 15:04:05. + zh_Hans: 日程开始时间,格式:2006-01-02 15:04:05。 + llm_description: 日程开始时间,格式:2006-01-02 15:04:05。 + form: llm + + - name: end_time + type: string + required: false + label: + en_US: End Time + zh_Hans: 结束时间 + human_description: + en_US: | + The end time of the event, format: 2006-01-02 15:04:05. + zh_Hans: 日程结束时间,格式:2006-01-02 15:04:05。 + llm_description: 日程结束时间,格式:2006-01-02 15:04:05。 + form: llm + + - name: auto_record + type: boolean + required: false + label: + en_US: Auto Record + zh_Hans: 自动录制 + human_description: + en_US: | + Whether to enable automatic recording, true: enabled, automatically record when the meeting starts; false: not enabled. + zh_Hans: 是否开启自动录制,true:开启,会议开始后自动录制;false:不开启。 + llm_description: 是否开启自动录制,true:开启,会议开始后自动录制;false:不开启。 + form: form diff --git a/api/core/tools/provider/builtin/lark_document/_assets/icon.svg b/api/core/tools/provider/builtin/lark_document/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..5a0a6416b3db3205b2e8c5d7039af120cfdd5b07 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_document/_assets/icon.svg @@ -0,0 +1,9 @@ + + + + + + + + + diff --git a/api/core/tools/provider/builtin/lark_document/lark_document.py b/api/core/tools/provider/builtin/lark_document/lark_document.py new file mode 100644 index 0000000000000000000000000000000000000000..b12832760283610c63ee81b2dd3a9da868ed55ec --- /dev/null +++ b/api/core/tools/provider/builtin/lark_document/lark_document.py @@ -0,0 +1,7 @@ +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController +from core.tools.utils.lark_api_utils import lark_auth + + +class LarkDocumentProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + lark_auth(credentials) diff --git a/api/core/tools/provider/builtin/lark_document/lark_document.yaml b/api/core/tools/provider/builtin/lark_document/lark_document.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0cb4ae1d62d3f866e2e8bcc4010cc01c841028ca --- /dev/null +++ b/api/core/tools/provider/builtin/lark_document/lark_document.yaml @@ -0,0 +1,36 @@ +identity: + author: Doug Lea + name: lark_document + label: + en_US: Lark Cloud Document + zh_Hans: Lark 云文档 + description: + en_US: | + Lark cloud document, requires the following permissions: docx:document、drive:drive、docs:document.content:read. + zh_Hans: | + Lark 云文档,需要开通以下权限: docx:document、drive:drive、docs:document.content:read。 + icon: icon.svg + tags: + - social + - productivity +credentials_for_provider: + app_id: + type: text-input + required: true + label: + en_US: APP ID + placeholder: + en_US: Please input your Lark app id + zh_Hans: 请输入你的 Lark app id + help: + en_US: Get your app_id and app_secret from Lark + zh_Hans: 从 Lark 获取您的 app_id 和 app_secret + url: https://open.larksuite.com/app + app_secret: + type: secret-input + required: true + label: + en_US: APP Secret + placeholder: + en_US: Please input your app secret + zh_Hans: 请输入你的 Lark app secret diff --git a/api/core/tools/provider/builtin/lark_document/tools/create_document.py b/api/core/tools/provider/builtin/lark_document/tools/create_document.py new file mode 100644 index 0000000000000000000000000000000000000000..2b1dae0db5578c1658db95e1494d25e09b103bb9 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_document/tools/create_document.py @@ -0,0 +1,19 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class CreateDocumentTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + title = tool_parameters.get("title") + content = tool_parameters.get("content") + folder_token = tool_parameters.get("folder_token") + + res = client.create_document(title, content, folder_token) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_document/tools/create_document.yaml b/api/core/tools/provider/builtin/lark_document/tools/create_document.yaml new file mode 100644 index 0000000000000000000000000000000000000000..37a1e23041c6c94e6a836b5d2e434188c3af3f80 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_document/tools/create_document.yaml @@ -0,0 +1,48 @@ +identity: + name: create_document + author: Doug Lea + label: + en_US: Create Lark document + zh_Hans: 创建 Lark 文档 +description: + human: + en_US: Create Lark document + zh_Hans: 创建 Lark 文档,支持创建空文档和带内容的文档,支持 markdown 语法创建。应用需要开启机器人能力(https://open.larksuite.com/document/faq/trouble-shooting/how-to-enable-bot-ability)。 + llm: A tool for creating Lark documents. +parameters: + - name: title + type: string + required: false + label: + en_US: Document title + zh_Hans: 文档标题 + human_description: + en_US: Document title, only supports plain text content. + zh_Hans: 文档标题,只支持纯文本内容。 + llm_description: 文档标题,只支持纯文本内容,可以为空。 + form: llm + + - name: content + type: string + required: false + label: + en_US: Document content + zh_Hans: 文档内容 + human_description: + en_US: Document content, supports markdown syntax, can be empty. + zh_Hans: 文档内容,支持 markdown 语法,可以为空。 + llm_description: 文档内容,支持 markdown 语法,可以为空。 + form: llm + + - name: folder_token + type: string + required: false + label: + en_US: folder_token + zh_Hans: 文档所在文件夹的 Token + human_description: + en_US: | + The token of the folder where the document is located. If it is not passed or is empty, it means the root directory. For Example: https://lark-japan.jp.larksuite.com/drive/folder/Lf8uf6BoAlWkUfdGtpMjUV0PpZd + zh_Hans: 文档所在文件夹的 Token,不传或传空表示根目录。例如:https://lark-japan.jp.larksuite.com/drive/folder/Lf8uf6BoAlWkUfdGtpMjUV0PpZd。 + llm_description: 文档所在文件夹的 Token,不传或传空表示根目录。例如:https://lark-japan.jp.larksuite.com/drive/folder/Lf8uf6BoAlWkUfdGtpMjUV0PpZd。 + form: llm diff --git a/api/core/tools/provider/builtin/lark_document/tools/get_document_content.py b/api/core/tools/provider/builtin/lark_document/tools/get_document_content.py new file mode 100644 index 0000000000000000000000000000000000000000..d15211b57e7a76e227f546e1262bbdba052c3bb5 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_document/tools/get_document_content.py @@ -0,0 +1,19 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class GetDocumentRawContentTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + document_id = tool_parameters.get("document_id") + mode = tool_parameters.get("mode", "markdown") + lang = tool_parameters.get("lang", "0") + + res = client.get_document_content(document_id, mode, lang) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_document/tools/get_document_content.yaml b/api/core/tools/provider/builtin/lark_document/tools/get_document_content.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fd6a033bfd6947431d82b20860134621811e0b96 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_document/tools/get_document_content.yaml @@ -0,0 +1,70 @@ +identity: + name: get_document_content + author: Doug Lea + label: + en_US: Get Lark Cloud Document Content + zh_Hans: 获取 Lark 云文档的内容 +description: + human: + en_US: Get lark cloud document content + zh_Hans: 获取 Lark 云文档的内容 + llm: A tool for retrieving content from Lark cloud documents. +parameters: + - name: document_id + type: string + required: true + label: + en_US: document_id + zh_Hans: Lark 文档的唯一标识 + human_description: + en_US: Unique identifier for a Lark document. You can also input the document's URL. + zh_Hans: Lark 文档的唯一标识,支持输入文档的 URL。 + llm_description: Lark 文档的唯一标识,支持输入文档的 URL。 + form: llm + + - name: mode + type: select + required: false + options: + - value: text + label: + en_US: text + zh_Hans: text + - value: markdown + label: + en_US: markdown + zh_Hans: markdown + default: "markdown" + label: + en_US: mode + zh_Hans: 文档返回格式 + human_description: + en_US: Format of the document return, optional values are text, markdown, can be empty, default is markdown. + zh_Hans: 文档返回格式,可选值有 text、markdown,可以为空,默认值为 markdown。 + llm_description: 文档返回格式,可选值有 text、markdown,可以为空,默认值为 markdown。 + form: form + + - name: lang + type: select + required: false + options: + - value: "0" + label: + en_US: User's default name + zh_Hans: 用户的默认名称 + - value: "1" + label: + en_US: User's English name + zh_Hans: 用户的英文名称 + default: "0" + label: + en_US: lang + zh_Hans: 指定@用户的语言 + human_description: + en_US: | + Specifies the language for MentionUser, optional values are [0, 1]. 0: User's default name, 1: User's English name, default is 0. + zh_Hans: | + 指定返回的 MentionUser,即@用户的语言,可选值有 [0,1]。0: 该用户的默认名称,1: 该用户的英文名称,默认值为 0。 + llm_description: | + 指定返回的 MentionUser,即@用户的语言,可选值有 [0,1]。0: 该用户的默认名称,1: 该用户的英文名称,默认值为 0。 + form: form diff --git a/api/core/tools/provider/builtin/lark_document/tools/list_document_blocks.py b/api/core/tools/provider/builtin/lark_document/tools/list_document_blocks.py new file mode 100644 index 0000000000000000000000000000000000000000..b96a87489e055e97e95129e98db00f512d27d4f6 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_document/tools/list_document_blocks.py @@ -0,0 +1,20 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class ListDocumentBlockTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + document_id = tool_parameters.get("document_id") + page_token = tool_parameters.get("page_token", "") + user_id_type = tool_parameters.get("user_id_type", "open_id") + page_size = tool_parameters.get("page_size", 500) + + res = client.list_document_blocks(document_id, page_token, user_id_type, page_size) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_document/tools/list_document_blocks.yaml b/api/core/tools/provider/builtin/lark_document/tools/list_document_blocks.yaml new file mode 100644 index 0000000000000000000000000000000000000000..08b673e0ae3ddcec733c4500d6f60beb56cd5326 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_document/tools/list_document_blocks.yaml @@ -0,0 +1,74 @@ +identity: + name: list_document_blocks + author: Doug Lea + label: + en_US: List Lark Document Blocks + zh_Hans: 获取 Lark 文档所有块 +description: + human: + en_US: List lark document blocks + zh_Hans: 获取 Lark 文档所有块的富文本内容并分页返回 + llm: A tool to get all blocks of Lark documents +parameters: + - name: document_id + type: string + required: true + label: + en_US: document_id + zh_Hans: Lark 文档的唯一标识 + human_description: + en_US: Unique identifier for a Lark document. You can also input the document's URL. + zh_Hans: Lark 文档的唯一标识,支持输入文档的 URL。 + llm_description: Lark 文档的唯一标识,支持输入文档的 URL。 + form: llm + + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form + + - name: page_size + type: number + required: false + default: 500 + label: + en_US: page_size + zh_Hans: 分页大小 + human_description: + en_US: Paging size, the default and maximum value is 500. + zh_Hans: 分页大小, 默认值和最大值为 500。 + llm_description: 分页大小, 表示一次请求最多返回多少条数据,默认值和最大值为 500。 + form: form + + - name: page_token + type: string + required: false + label: + en_US: page_token + zh_Hans: 分页标记 + human_description: + en_US: Pagination token used to navigate through query results, allowing retrieval of additional items in subsequent requests. + zh_Hans: 分页标记,用于分页查询结果,以便下次遍历时获取更多项。 + llm_description: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + form: llm diff --git a/api/core/tools/provider/builtin/lark_document/tools/write_document.py b/api/core/tools/provider/builtin/lark_document/tools/write_document.py new file mode 100644 index 0000000000000000000000000000000000000000..888e0e39fce3894e8a62052172f217087da8917d --- /dev/null +++ b/api/core/tools/provider/builtin/lark_document/tools/write_document.py @@ -0,0 +1,19 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class CreateDocumentTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + document_id = tool_parameters.get("document_id") + content = tool_parameters.get("content") + position = tool_parameters.get("position", "end") + + res = client.write_document(document_id, content, position) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_document/tools/write_document.yaml b/api/core/tools/provider/builtin/lark_document/tools/write_document.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9cdf034ed082304d0a1e2e2648924ab7e863a2c3 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_document/tools/write_document.yaml @@ -0,0 +1,57 @@ +identity: + name: write_document + author: Doug Lea + label: + en_US: Write Document + zh_Hans: 在 Lark 文档中新增内容 +description: + human: + en_US: Adding new content to Lark documents + zh_Hans: 在 Lark 文档中新增内容 + llm: A tool for adding new content to Lark documents. +parameters: + - name: document_id + type: string + required: true + label: + en_US: document_id + zh_Hans: Lark 文档的唯一标识 + human_description: + en_US: Unique identifier for a Lark document. You can also input the document's URL. + zh_Hans: Lark 文档的唯一标识,支持输入文档的 URL。 + llm_description: Lark 文档的唯一标识,支持输入文档的 URL。 + form: llm + + - name: content + type: string + required: true + label: + en_US: Plain text or Markdown content + zh_Hans: 纯文本或 Markdown 内容 + human_description: + en_US: Plain text or Markdown content. Note that embedded tables in the document should not have merged cells. + zh_Hans: 纯文本或 Markdown 内容。注意文档的内嵌套表格不允许有单元格合并。 + llm_description: 纯文本或 Markdown 内容,注意文档的内嵌套表格不允许有单元格合并。 + form: llm + + - name: position + type: select + required: false + options: + - value: start + label: + en_US: document start + zh_Hans: 文档开始 + - value: end + label: + en_US: document end + zh_Hans: 文档结束 + default: "end" + label: + en_US: position + zh_Hans: 内容添加位置 + human_description: + en_US: Content insertion position, optional values are start, end. 'start' means adding content at the beginning of the document; 'end' means adding content at the end of the document. The default value is end. + zh_Hans: 内容添加位置,可选值有 start、end。start 表示在文档开头添加内容;end 表示在文档结尾添加内容,默认值为 end。 + llm_description: 内容添加位置,可选值有 start、end。start 表示在文档开头添加内容;end 表示在文档结尾添加内容,默认值为 end。 + form: form diff --git a/api/core/tools/provider/builtin/lark_message_and_group/_assets/icon.png b/api/core/tools/provider/builtin/lark_message_and_group/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..0dfd58a9d512fd0f0481e32f913ba978d6219002 Binary files /dev/null and b/api/core/tools/provider/builtin/lark_message_and_group/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/lark_message_and_group/lark_message_and_group.py b/api/core/tools/provider/builtin/lark_message_and_group/lark_message_and_group.py new file mode 100644 index 0000000000000000000000000000000000000000..de6997b0bf942f5fd02b618424f60fb1d4adb614 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_message_and_group/lark_message_and_group.py @@ -0,0 +1,7 @@ +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController +from core.tools.utils.lark_api_utils import lark_auth + + +class LarkMessageAndGroupProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + lark_auth(credentials) diff --git a/api/core/tools/provider/builtin/lark_message_and_group/lark_message_and_group.yaml b/api/core/tools/provider/builtin/lark_message_and_group/lark_message_and_group.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ad3fe0f36190989ad8e7353acd9f1b8896e3ed3a --- /dev/null +++ b/api/core/tools/provider/builtin/lark_message_and_group/lark_message_and_group.yaml @@ -0,0 +1,36 @@ +identity: + author: Doug Lea + name: lark_message_and_group + label: + en_US: Lark Message And Group + zh_Hans: Lark 消息和群组 + description: + en_US: | + Lark message and group, requires the following permissions: im:message、im:message.group_msg. + zh_Hans: | + Lark 消息和群组,需要开通以下权限: im:message、im:message.group_msg。 + icon: icon.png + tags: + - social + - productivity +credentials_for_provider: + app_id: + type: text-input + required: true + label: + en_US: APP ID + placeholder: + en_US: Please input your Lark app id + zh_Hans: 请输入你的 Lark app id + help: + en_US: Get your app_id and app_secret from Lark + zh_Hans: 从 Lark 获取您的 app_id 和 app_secret + url: https://open.larksuite.com/app + app_secret: + type: secret-input + required: true + label: + en_US: APP Secret + placeholder: + en_US: Please input your app secret + zh_Hans: 请输入你的 Lark app secret diff --git a/api/core/tools/provider/builtin/lark_message_and_group/tools/get_chat_messages.py b/api/core/tools/provider/builtin/lark_message_and_group/tools/get_chat_messages.py new file mode 100644 index 0000000000000000000000000000000000000000..118bac7ab7d720f242b1f6164eae59ed89dffc30 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_message_and_group/tools/get_chat_messages.py @@ -0,0 +1,23 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class GetChatMessagesTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + container_id = tool_parameters.get("container_id") + start_time = tool_parameters.get("start_time") + end_time = tool_parameters.get("end_time") + page_token = tool_parameters.get("page_token") + sort_type = tool_parameters.get("sort_type", "ByCreateTimeAsc") + page_size = tool_parameters.get("page_size", 20) + + res = client.get_chat_messages(container_id, start_time, end_time, page_token, sort_type, page_size) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_message_and_group/tools/get_chat_messages.yaml b/api/core/tools/provider/builtin/lark_message_and_group/tools/get_chat_messages.yaml new file mode 100644 index 0000000000000000000000000000000000000000..965b45a5fbaec9050e10da79f2c7485b6963192f --- /dev/null +++ b/api/core/tools/provider/builtin/lark_message_and_group/tools/get_chat_messages.yaml @@ -0,0 +1,96 @@ +identity: + name: get_chat_messages + author: Doug Lea + label: + en_US: Get Chat Messages + zh_Hans: 获取指定单聊、群聊的消息历史 +description: + human: + en_US: Get Chat Messages + zh_Hans: 获取指定单聊、群聊的消息历史 + llm: A tool for getting chat messages from specific one-on-one chats or group chats.(获取指定单聊、群聊的消息历史) +parameters: + - name: container_id + type: string + required: true + label: + en_US: Container Id + zh_Hans: 群聊或单聊的 ID + human_description: + en_US: The ID of the group chat or single chat. Refer to the group ID description for how to obtain it. https://open.larkoffice.com/document/server-docs/group/chat/chat-id-description + zh_Hans: 群聊或单聊的 ID,获取方式参见群 ID 说明。https://open.larkoffice.com/document/server-docs/group/chat/chat-id-description + llm_description: 群聊或单聊的 ID,获取方式参见群 ID 说明。https://open.larkoffice.com/document/server-docs/group/chat/chat-id-description + form: llm + + - name: start_time + type: string + required: false + label: + en_US: Start Time + zh_Hans: 起始时间 + human_description: + en_US: The start time for querying historical messages, formatted as "2006-01-02 15:04:05". + zh_Hans: 待查询历史信息的起始时间,格式为 "2006-01-02 15:04:05"。 + llm_description: 待查询历史信息的起始时间,格式为 "2006-01-02 15:04:05"。 + form: llm + + - name: end_time + type: string + required: false + label: + en_US: End Time + zh_Hans: 结束时间 + human_description: + en_US: The end time for querying historical messages, formatted as "2006-01-02 15:04:05". + zh_Hans: 待查询历史信息的结束时间,格式为 "2006-01-02 15:04:05"。 + llm_description: 待查询历史信息的结束时间,格式为 "2006-01-02 15:04:05"。 + form: llm + + - name: sort_type + type: select + required: false + options: + - value: ByCreateTimeAsc + label: + en_US: ByCreateTimeAsc + zh_Hans: ByCreateTimeAsc + - value: ByCreateTimeDesc + label: + en_US: ByCreateTimeDesc + zh_Hans: ByCreateTimeDesc + default: "ByCreateTimeAsc" + label: + en_US: Sort Type + zh_Hans: 排序方式 + human_description: + en_US: | + The message sorting method. Optional values are ByCreateTimeAsc: sorted in ascending order by message creation time; ByCreateTimeDesc: sorted in descending order by message creation time. The default value is ByCreateTimeAsc. Note: When using page_token for pagination requests, the sorting method (sort_type) is consistent with the first request and cannot be changed midway. + zh_Hans: | + 消息排序方式,可选值有 ByCreateTimeAsc:按消息创建时间升序排列;ByCreateTimeDesc:按消息创建时间降序排列。默认值为:ByCreateTimeAsc。注意:使用 page_token 分页请求时,排序方式(sort_type)均与第一次请求一致,不支持中途改换排序方式。 + llm_description: 消息排序方式,可选值有 ByCreateTimeAsc:按消息创建时间升序排列;ByCreateTimeDesc:按消息创建时间降序排列。默认值为:ByCreateTimeAsc。注意:使用 page_token 分页请求时,排序方式(sort_type)均与第一次请求一致,不支持中途改换排序方式。 + form: form + + - name: page_size + type: number + required: false + default: 20 + label: + en_US: Page Size + zh_Hans: 分页大小 + human_description: + en_US: The page size, i.e., the number of data entries returned in a single request. The default value is 20, and the value range is [1,50]. + zh_Hans: 分页大小,即单次请求所返回的数据条目数。默认值为 20,取值范围为 [1,50]。 + llm_description: 分页大小,即单次请求所返回的数据条目数。默认值为 20,取值范围为 [1,50]。 + form: form + + - name: page_token + type: string + required: false + label: + en_US: Page Token + zh_Hans: 分页标记 + human_description: + en_US: The pagination token. Leave it blank for the first request, indicating to start traversing from the beginning; when the pagination query result has more items, a new page_token will be returned simultaneously, which can be used to obtain the query result in the next traversal. + zh_Hans: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + llm_description: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + form: llm diff --git a/api/core/tools/provider/builtin/lark_message_and_group/tools/get_thread_messages.py b/api/core/tools/provider/builtin/lark_message_and_group/tools/get_thread_messages.py new file mode 100644 index 0000000000000000000000000000000000000000..3509d9bbcfe437088cb0e7ce148ce20df40cb1df --- /dev/null +++ b/api/core/tools/provider/builtin/lark_message_and_group/tools/get_thread_messages.py @@ -0,0 +1,21 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class GetChatMessagesTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + container_id = tool_parameters.get("container_id") + page_token = tool_parameters.get("page_token") + sort_type = tool_parameters.get("sort_type", "ByCreateTimeAsc") + page_size = tool_parameters.get("page_size", 20) + + res = client.get_thread_messages(container_id, page_token, sort_type, page_size) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_message_and_group/tools/get_thread_messages.yaml b/api/core/tools/provider/builtin/lark_message_and_group/tools/get_thread_messages.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5f7a4f0902523e49775c4ff694cad1c2d2e34cb4 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_message_and_group/tools/get_thread_messages.yaml @@ -0,0 +1,72 @@ +identity: + name: get_thread_messages + author: Doug Lea + label: + en_US: Get Thread Messages + zh_Hans: 获取指定话题的消息历史 +description: + human: + en_US: Get Thread Messages + zh_Hans: 获取指定话题的消息历史 + llm: A tool for getting chat messages from specific threads.(获取指定话题的消息历史) +parameters: + - name: container_id + type: string + required: true + label: + en_US: Thread Id + zh_Hans: 话题 ID + human_description: + en_US: The ID of the thread. Refer to the thread overview on how to obtain the thread_id. https://open.larksuite.com/document/uAjLw4CM/ukTMukTMukTM/reference/im-v1/message/thread-introduction + zh_Hans: 话题 ID,获取方式参见话题概述的如何获取 thread_id 章节。https://open.larksuite.com/document/uAjLw4CM/ukTMukTMukTM/reference/im-v1/message/thread-introduction + llm_description: 话题 ID,获取方式参见话题概述的如何获取 thread_id 章节。https://open.larksuite.com/document/uAjLw4CM/ukTMukTMukTM/reference/im-v1/message/thread-introduction + form: llm + + - name: sort_type + type: select + required: false + options: + - value: ByCreateTimeAsc + label: + en_US: ByCreateTimeAsc + zh_Hans: ByCreateTimeAsc + - value: ByCreateTimeDesc + label: + en_US: ByCreateTimeDesc + zh_Hans: ByCreateTimeDesc + default: "ByCreateTimeAsc" + label: + en_US: Sort Type + zh_Hans: 排序方式 + human_description: + en_US: | + The message sorting method. Optional values are ByCreateTimeAsc: sorted in ascending order by message creation time; ByCreateTimeDesc: sorted in descending order by message creation time. The default value is ByCreateTimeAsc. Note: When using page_token for pagination requests, the sorting method (sort_type) is consistent with the first request and cannot be changed midway. + zh_Hans: | + 消息排序方式,可选值有 ByCreateTimeAsc:按消息创建时间升序排列;ByCreateTimeDesc:按消息创建时间降序排列。默认值为:ByCreateTimeAsc。注意:使用 page_token 分页请求时,排序方式(sort_type)均与第一次请求一致,不支持中途改换排序方式。 + llm_description: 消息排序方式,可选值有 ByCreateTimeAsc:按消息创建时间升序排列;ByCreateTimeDesc:按消息创建时间降序排列。默认值为:ByCreateTimeAsc。注意:使用 page_token 分页请求时,排序方式(sort_type)均与第一次请求一致,不支持中途改换排序方式。 + form: form + + - name: page_size + type: number + required: false + default: 20 + label: + en_US: Page Size + zh_Hans: 分页大小 + human_description: + en_US: The page size, i.e., the number of data entries returned in a single request. The default value is 20, and the value range is [1,50]. + zh_Hans: 分页大小,即单次请求所返回的数据条目数。默认值为 20,取值范围为 [1,50]。 + llm_description: 分页大小,即单次请求所返回的数据条目数。默认值为 20,取值范围为 [1,50]。 + form: form + + - name: page_token + type: string + required: false + label: + en_US: Page Token + zh_Hans: 分页标记 + human_description: + en_US: The pagination token. Leave it blank for the first request, indicating to start traversing from the beginning; when the pagination query result has more items, a new page_token will be returned simultaneously, which can be used to obtain the query result in the next traversal. + zh_Hans: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + llm_description: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + form: llm diff --git a/api/core/tools/provider/builtin/lark_message_and_group/tools/send_bot_message.py b/api/core/tools/provider/builtin/lark_message_and_group/tools/send_bot_message.py new file mode 100644 index 0000000000000000000000000000000000000000..b0a8df61e85f2eb73e86473ee9de8a661d1cf046 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_message_and_group/tools/send_bot_message.py @@ -0,0 +1,20 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class SendBotMessageTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + receive_id_type = tool_parameters.get("receive_id_type") + receive_id = tool_parameters.get("receive_id") + msg_type = tool_parameters.get("msg_type") + content = tool_parameters.get("content") + + res = client.send_bot_message(receive_id_type, receive_id, msg_type, content) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_message_and_group/tools/send_bot_message.yaml b/api/core/tools/provider/builtin/lark_message_and_group/tools/send_bot_message.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b949c5e01694ce5ff00dbec07b12597be0315b08 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_message_and_group/tools/send_bot_message.yaml @@ -0,0 +1,125 @@ +identity: + name: send_bot_message + author: Doug Lea + label: + en_US: Send Bot Message + zh_Hans: 发送 Lark 应用消息 +description: + human: + en_US: Send bot message + zh_Hans: 发送 Lark 应用消息 + llm: A tool for sending Lark application messages. +parameters: + - name: receive_id + type: string + required: true + label: + en_US: receive_id + zh_Hans: 消息接收者的 ID + human_description: + en_US: The ID of the message receiver, the ID type is consistent with the value of the query parameter receive_id_type. + zh_Hans: 消息接收者的 ID,ID 类型与查询参数 receive_id_type 的取值一致。 + llm_description: 消息接收者的 ID,ID 类型与查询参数 receive_id_type 的取值一致。 + form: llm + + - name: receive_id_type + type: select + required: true + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + - value: email + label: + en_US: email + zh_Hans: email + - value: chat_id + label: + en_US: chat_id + zh_Hans: chat_id + label: + en_US: receive_id_type + zh_Hans: 消息接收者的 ID 类型 + human_description: + en_US: The ID type of the message receiver, optional values are open_id, union_id, user_id, email, chat_id, with a default value of open_id. + zh_Hans: 消息接收者的 ID 类型,可选值有 open_id、union_id、user_id、email、chat_id,默认值为 open_id。 + llm_description: 消息接收者的 ID 类型,可选值有 open_id、union_id、user_id、email、chat_id,默认值为 open_id。 + form: form + + - name: msg_type + type: select + required: true + options: + - value: text + label: + en_US: text + zh_Hans: 文本 + - value: interactive + label: + en_US: interactive + zh_Hans: 卡片 + - value: post + label: + en_US: post + zh_Hans: 富文本 + - value: image + label: + en_US: image + zh_Hans: 图片 + - value: file + label: + en_US: file + zh_Hans: 文件 + - value: audio + label: + en_US: audio + zh_Hans: 语音 + - value: media + label: + en_US: media + zh_Hans: 视频 + - value: sticker + label: + en_US: sticker + zh_Hans: 表情包 + - value: share_chat + label: + en_US: share_chat + zh_Hans: 分享群名片 + - value: share_user + label: + en_US: share_user + zh_Hans: 分享个人名片 + - value: system + label: + en_US: system + zh_Hans: 系统消息 + label: + en_US: msg_type + zh_Hans: 消息类型 + human_description: + en_US: Message type. Optional values are text, post, image, file, audio, media, sticker, interactive, share_chat, share_user, system. For detailed introduction of different message types, refer to the message content(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json). + zh_Hans: 消息类型。可选值有:text、post、image、file、audio、media、sticker、interactive、share_chat、share_user、system。不同消息类型的详细介绍,参见发送消息内容(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json)。 + llm_description: 消息类型。可选值有:text、post、image、file、audio、media、sticker、interactive、share_chat、share_user、system。不同消息类型的详细介绍,参见发送消息内容(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json)。 + form: form + + - name: content + type: string + required: true + label: + en_US: content + zh_Hans: 消息内容 + human_description: + en_US: Message content, a JSON structure serialized string. The value of this parameter corresponds to msg_type. For example, if msg_type is text, this parameter needs to pass in text type content. To understand the format and usage limitations of different message types, refer to the message content(https://open.larksuite.com/document/server-docs/im-v1/message-content-description/create_json). + zh_Hans: 消息内容,JSON 结构序列化后的字符串。该参数的取值与 msg_type 对应,例如 msg_type 取值为 text,则该参数需要传入文本类型的内容。了解不同类型的消息内容格式、使用限制,可参见发送消息内容(https://open.larksuite.com/document/server-docs/im-v1/message-content-description/create_json)。 + llm_description: 消息内容,JSON 结构序列化后的字符串。该参数的取值与 msg_type 对应,例如 msg_type 取值为 text,则该参数需要传入文本类型的内容。了解不同类型的消息内容格式、使用限制,可参见发送消息内容(https://open.larksuite.com/document/server-docs/im-v1/message-content-description/create_json)。 + form: llm diff --git a/api/core/tools/provider/builtin/lark_message_and_group/tools/send_webhook_message.py b/api/core/tools/provider/builtin/lark_message_and_group/tools/send_webhook_message.py new file mode 100644 index 0000000000000000000000000000000000000000..18a605079fc9509e06a2f53d90005f9303f7df4a --- /dev/null +++ b/api/core/tools/provider/builtin/lark_message_and_group/tools/send_webhook_message.py @@ -0,0 +1,19 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class SendWebhookMessageTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + webhook = tool_parameters.get("webhook") + msg_type = tool_parameters.get("msg_type") + content = tool_parameters.get("content") + + res = client.send_webhook_message(webhook, msg_type, content) + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_message_and_group/tools/send_webhook_message.yaml b/api/core/tools/provider/builtin/lark_message_and_group/tools/send_webhook_message.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ea13cae52ba997fe05b92f3a5f5298fc2b23838e --- /dev/null +++ b/api/core/tools/provider/builtin/lark_message_and_group/tools/send_webhook_message.yaml @@ -0,0 +1,68 @@ +identity: + name: send_webhook_message + author: Doug Lea + label: + en_US: Send Webhook Message + zh_Hans: 使用自定义机器人发送 Lark 消息 +description: + human: + en_US: Send webhook message + zh_Hans: 使用自定义机器人发送 Lark 消息 + llm: A tool for sending Lark messages using a custom robot. +parameters: + - name: webhook + type: string + required: true + label: + en_US: webhook + zh_Hans: webhook + human_description: + en_US: | + The address of the webhook, the format of the webhook address corresponding to the bot is as follows: https://open.larksuite.com/open-apis/bot/v2/hook/xxxxxxxxxxxxxxxxx. For details, please refer to: Lark Custom Bot Usage Guide(https://open.larkoffice.com/document/client-docs/bot-v3/add-custom-bot) + zh_Hans: | + webhook 的地址,机器人对应的 webhook 地址格式如下: https://open.larksuite.com/open-apis/bot/v2/hook/xxxxxxxxxxxxxxxxx,详情可参考: Lark 自定义机器人使用指南(https://open.larksuite.com/document/client-docs/bot-v3/add-custom-bot) + llm_description: | + webhook 的地址,机器人对应的 webhook 地址格式如下: https://open.larksuite.com/open-apis/bot/v2/hook/xxxxxxxxxxxxxxxxx,详情可参考: Lark 自定义机器人使用指南(https://open.larksuite.com/document/client-docs/bot-v3/add-custom-bot) + form: llm + + - name: msg_type + type: select + required: true + options: + - value: text + label: + en_US: text + zh_Hans: 文本 + - value: interactive + label: + en_US: interactive + zh_Hans: 卡片 + - value: image + label: + en_US: image + zh_Hans: 图片 + - value: share_chat + label: + en_US: share_chat + zh_Hans: 分享群名片 + label: + en_US: msg_type + zh_Hans: 消息类型 + human_description: + en_US: Message type. Optional values are text, image, interactive, share_chat. For detailed introduction of different message types, refer to the message content(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json). + zh_Hans: 消息类型。可选值有:text、image、interactive、share_chat。不同消息类型的详细介绍,参见发送消息内容(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json)。 + llm_description: 消息类型。可选值有:text、image、interactive、share_chat。不同消息类型的详细介绍,参见发送消息内容(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json)。 + form: form + + + - name: content + type: string + required: true + label: + en_US: content + zh_Hans: 消息内容 + human_description: + en_US: Message content, a JSON structure serialized string. The value of this parameter corresponds to msg_type. For example, if msg_type is text, this parameter needs to pass in text type content. To understand the format and usage limitations of different message types, refer to the message content(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json). + zh_Hans: 消息内容,JSON 结构序列化后的字符串。该参数的取值与 msg_type 对应,例如 msg_type 取值为 text,则该参数需要传入文本类型的内容。了解不同类型的消息内容格式、使用限制,可参见发送消息内容(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json)。 + llm_description: 消息内容,JSON 结构序列化后的字符串。该参数的取值与 msg_type 对应,例如 msg_type 取值为 text,则该参数需要传入文本类型的内容。了解不同类型的消息内容格式、使用限制,可参见发送消息内容(https://open.larkoffice.com/document/server-docs/im-v1/message-content-description/create_json)。 + form: llm diff --git a/api/core/tools/provider/builtin/lark_spreadsheet/_assets/icon.png b/api/core/tools/provider/builtin/lark_spreadsheet/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..258b361261d4e3366251613141efaf200cd492db Binary files /dev/null and b/api/core/tools/provider/builtin/lark_spreadsheet/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/lark_spreadsheet/lark_spreadsheet.py b/api/core/tools/provider/builtin/lark_spreadsheet/lark_spreadsheet.py new file mode 100644 index 0000000000000000000000000000000000000000..c791363f21fbe12f9116583adcfdf84a5d329b25 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_spreadsheet/lark_spreadsheet.py @@ -0,0 +1,7 @@ +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController +from core.tools.utils.lark_api_utils import lark_auth + + +class LarkMessageProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + lark_auth(credentials) diff --git a/api/core/tools/provider/builtin/lark_spreadsheet/lark_spreadsheet.yaml b/api/core/tools/provider/builtin/lark_spreadsheet/lark_spreadsheet.yaml new file mode 100644 index 0000000000000000000000000000000000000000..030b5c9063227acacacef8b879e0f52ae0211517 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_spreadsheet/lark_spreadsheet.yaml @@ -0,0 +1,36 @@ +identity: + author: Doug Lea + name: lark_spreadsheet + label: + en_US: Lark Spreadsheet + zh_Hans: Lark 电子表格 + description: + en_US: | + Lark Spreadsheet, requires the following permissions: sheets:spreadsheet. + zh_Hans: | + Lark 电子表格,需要开通以下权限: sheets:spreadsheet。 + icon: icon.png + tags: + - social + - productivity +credentials_for_provider: + app_id: + type: text-input + required: true + label: + en_US: APP ID + placeholder: + en_US: Please input your Lark app id + zh_Hans: 请输入你的 Lark app id + help: + en_US: Get your app_id and app_secret from Lark + zh_Hans: 从 Lark 获取您的 app_id 和 app_secret + url: https://open.larksuite.com/app + app_secret: + type: secret-input + required: true + label: + en_US: APP Secret + placeholder: + en_US: Please input your app secret + zh_Hans: 请输入你的 Lark app secret diff --git a/api/core/tools/provider/builtin/lark_spreadsheet/tools/add_cols.py b/api/core/tools/provider/builtin/lark_spreadsheet/tools/add_cols.py new file mode 100644 index 0000000000000000000000000000000000000000..deeb5a1ecf6f7d38dc32523074dad3b71593376a --- /dev/null +++ b/api/core/tools/provider/builtin/lark_spreadsheet/tools/add_cols.py @@ -0,0 +1,22 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class AddColsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + spreadsheet_token = tool_parameters.get("spreadsheet_token") + sheet_id = tool_parameters.get("sheet_id") + sheet_name = tool_parameters.get("sheet_name") + length = tool_parameters.get("length") + values = tool_parameters.get("values") + + res = client.add_cols(spreadsheet_token, sheet_id, sheet_name, length, values) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_spreadsheet/tools/add_cols.yaml b/api/core/tools/provider/builtin/lark_spreadsheet/tools/add_cols.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b73335f405c20c7ac3405552669ec34fc0cd4754 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_spreadsheet/tools/add_cols.yaml @@ -0,0 +1,72 @@ +identity: + name: add_cols + author: Doug Lea + label: + en_US: Add Cols + zh_Hans: 新增多列至工作表最后 +description: + human: + en_US: Add Cols + zh_Hans: 新增多列至工作表最后 + llm: A tool for adding multiple columns to the end of a spreadsheet. (新增多列至工作表最后) +parameters: + - name: spreadsheet_token + type: string + required: true + label: + en_US: spreadsheet_token + zh_Hans: 电子表格 token + human_description: + en_US: Spreadsheet token, supports input of spreadsheet URL. + zh_Hans: 电子表格 token,支持输入电子表格 url。 + llm_description: 电子表格 token,支持输入电子表格 url。 + form: llm + + - name: sheet_id + type: string + required: false + label: + en_US: sheet_id + zh_Hans: 工作表 ID + human_description: + en_US: Sheet ID, either sheet_id or sheet_name must be filled. + zh_Hans: 工作表 ID,与 sheet_name 二者其一必填。 + llm_description: 工作表 ID,与 sheet_name 二者其一必填。 + form: llm + + - name: sheet_name + type: string + required: false + label: + en_US: sheet_name + zh_Hans: 工作表名称 + human_description: + en_US: Sheet name, either sheet_id or sheet_name must be filled. + zh_Hans: 工作表名称,与 sheet_id 二者其一必填。 + llm_description: 工作表名称,与 sheet_id 二者其一必填。 + form: llm + + - name: length + type: number + required: true + label: + en_US: length + zh_Hans: 要增加的列数 + human_description: + en_US: Number of columns to add, range (0-5000]. + zh_Hans: 要增加的列数,范围(0-5000]。 + llm_description: 要增加的列数,范围(0-5000]。 + form: form + + - name: values + type: string + required: false + label: + en_US: values + zh_Hans: 新增列的单元格内容 + human_description: + en_US: | + Content of the new columns, array of objects in string format, each array represents a row of table data, format like: [ [ "ID","Name","Age" ],[ 1,"Zhang San",10 ],[ 2,"Li Si",11 ] ]. + zh_Hans: 新增列的单元格内容,数组对象字符串,每个数组一行表格数据,格式:[["编号","姓名","年龄"],[1,"张三",10],[2,"李四",11]]。 + llm_description: 新增列的单元格内容,数组对象字符串,每个数组一行表格数据,格式:[["编号","姓名","年龄"],[1,"张三",10],[2,"李四",11]]。 + form: llm diff --git a/api/core/tools/provider/builtin/lark_spreadsheet/tools/add_rows.py b/api/core/tools/provider/builtin/lark_spreadsheet/tools/add_rows.py new file mode 100644 index 0000000000000000000000000000000000000000..f434b1c60341f34a40e1c7f88c72f29ee230c396 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_spreadsheet/tools/add_rows.py @@ -0,0 +1,22 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class AddRowsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + spreadsheet_token = tool_parameters.get("spreadsheet_token") + sheet_id = tool_parameters.get("sheet_id") + sheet_name = tool_parameters.get("sheet_name") + length = tool_parameters.get("length") + values = tool_parameters.get("values") + + res = client.add_rows(spreadsheet_token, sheet_id, sheet_name, length, values) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_spreadsheet/tools/add_rows.yaml b/api/core/tools/provider/builtin/lark_spreadsheet/tools/add_rows.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6bce305b9825ec8c826b54edbd94bbc040c75bcc --- /dev/null +++ b/api/core/tools/provider/builtin/lark_spreadsheet/tools/add_rows.yaml @@ -0,0 +1,72 @@ +identity: + name: add_rows + author: Doug Lea + label: + en_US: Add Rows + zh_Hans: 新增多行至工作表最后 +description: + human: + en_US: Add Rows + zh_Hans: 新增多行至工作表最后 + llm: A tool for adding multiple rows to the end of a spreadsheet. (新增多行至工作表最后) +parameters: + - name: spreadsheet_token + type: string + required: true + label: + en_US: spreadsheet_token + zh_Hans: 电子表格 token + human_description: + en_US: Spreadsheet token, supports input of spreadsheet URL. + zh_Hans: 电子表格 token,支持输入电子表格 url。 + llm_description: 电子表格 token,支持输入电子表格 url。 + form: llm + + - name: sheet_id + type: string + required: false + label: + en_US: sheet_id + zh_Hans: 工作表 ID + human_description: + en_US: Sheet ID, either sheet_id or sheet_name must be filled. + zh_Hans: 工作表 ID,与 sheet_name 二者其一必填。 + llm_description: 工作表 ID,与 sheet_name 二者其一必填。 + form: llm + + - name: sheet_name + type: string + required: false + label: + en_US: sheet_name + zh_Hans: 工作表名称 + human_description: + en_US: Sheet name, either sheet_id or sheet_name must be filled. + zh_Hans: 工作表名称,与 sheet_id 二者其一必填。 + llm_description: 工作表名称,与 sheet_id 二者其一必填。 + form: llm + + - name: length + type: number + required: true + label: + en_US: length + zh_Hans: 要增加行数 + human_description: + en_US: Number of rows to add, range (0-5000]. + zh_Hans: 要增加行数,范围(0-5000]。 + llm_description: 要增加行数,范围(0-5000]。 + form: form + + - name: values + type: string + required: false + label: + en_US: values + zh_Hans: 新增行的表格内容 + human_description: + en_US: | + Content of the new rows, array of objects in string format, each array represents a row of table data, format like: [ [ "ID","Name","Age" ],[ 1,"Zhang San",10 ],[ 2,"Li Si",11 ] ]. + zh_Hans: 新增行的表格内容,数组对象字符串,每个数组一行表格数据,格式,如:[["编号","姓名","年龄"],[1,"张三",10],[2,"李四",11]]。 + llm_description: 新增行的表格内容,数组对象字符串,每个数组一行表格数据,格式,如:[["编号","姓名","年龄"],[1,"张三",10],[2,"李四",11]]。 + form: llm diff --git a/api/core/tools/provider/builtin/lark_spreadsheet/tools/create_spreadsheet.py b/api/core/tools/provider/builtin/lark_spreadsheet/tools/create_spreadsheet.py new file mode 100644 index 0000000000000000000000000000000000000000..74b20ac2c838f8bd49404e89416e359042106c1e --- /dev/null +++ b/api/core/tools/provider/builtin/lark_spreadsheet/tools/create_spreadsheet.py @@ -0,0 +1,19 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class CreateSpreadsheetTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + title = tool_parameters.get("title") + folder_token = tool_parameters.get("folder_token") + + res = client.create_spreadsheet(title, folder_token) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_spreadsheet/tools/create_spreadsheet.yaml b/api/core/tools/provider/builtin/lark_spreadsheet/tools/create_spreadsheet.yaml new file mode 100644 index 0000000000000000000000000000000000000000..931310e63172d4227fd8663f66c68d201d73aa8f --- /dev/null +++ b/api/core/tools/provider/builtin/lark_spreadsheet/tools/create_spreadsheet.yaml @@ -0,0 +1,35 @@ +identity: + name: create_spreadsheet + author: Doug Lea + label: + en_US: Create Spreadsheet + zh_Hans: 创建电子表格 +description: + human: + en_US: Create Spreadsheet + zh_Hans: 创建电子表格 + llm: A tool for creating spreadsheets. (创建电子表格) +parameters: + - name: title + type: string + required: false + label: + en_US: Spreadsheet Title + zh_Hans: 电子表格标题 + human_description: + en_US: The title of the spreadsheet + zh_Hans: 电子表格的标题 + llm_description: 电子表格的标题 + form: llm + + - name: folder_token + type: string + required: false + label: + en_US: Folder Token + zh_Hans: 文件夹 token + human_description: + en_US: The token of the folder, supports folder URL input, e.g., https://bytedance.larkoffice.com/drive/folder/CxHEf4DCSlNkL2dUTCJcPRgentg + zh_Hans: 文件夹 token,支持文件夹 URL 输入,如:https://bytedance.larkoffice.com/drive/folder/CxHEf4DCSlNkL2dUTCJcPRgentg + llm_description: 文件夹 token,支持文件夹 URL 输入,如:https://bytedance.larkoffice.com/drive/folder/CxHEf4DCSlNkL2dUTCJcPRgentg + form: llm diff --git a/api/core/tools/provider/builtin/lark_spreadsheet/tools/get_spreadsheet.py b/api/core/tools/provider/builtin/lark_spreadsheet/tools/get_spreadsheet.py new file mode 100644 index 0000000000000000000000000000000000000000..0fe35b6dc645b86fba9d14c7b02884f4fde3a069 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_spreadsheet/tools/get_spreadsheet.py @@ -0,0 +1,19 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class GetSpreadsheetTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + spreadsheet_token = tool_parameters.get("spreadsheet_token") + user_id_type = tool_parameters.get("user_id_type", "open_id") + + res = client.get_spreadsheet(spreadsheet_token, user_id_type) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_spreadsheet/tools/get_spreadsheet.yaml b/api/core/tools/provider/builtin/lark_spreadsheet/tools/get_spreadsheet.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c519938617ba8c331467aa6bf6c00c283e9b42be --- /dev/null +++ b/api/core/tools/provider/builtin/lark_spreadsheet/tools/get_spreadsheet.yaml @@ -0,0 +1,49 @@ +identity: + name: get_spreadsheet + author: Doug Lea + label: + en_US: Get Spreadsheet + zh_Hans: 获取电子表格信息 +description: + human: + en_US: Get Spreadsheet + zh_Hans: 获取电子表格信息 + llm: A tool for getting information from spreadsheets. (获取电子表格信息) +parameters: + - name: spreadsheet_token + type: string + required: true + label: + en_US: Spreadsheet Token + zh_Hans: 电子表格 token + human_description: + en_US: Spreadsheet token, supports input of spreadsheet URL. + zh_Hans: 电子表格 token,支持输入电子表格 URL。 + llm_description: 电子表格 token,支持输入电子表格 URL。 + form: llm + + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form diff --git a/api/core/tools/provider/builtin/lark_spreadsheet/tools/list_spreadsheet_sheets.py b/api/core/tools/provider/builtin/lark_spreadsheet/tools/list_spreadsheet_sheets.py new file mode 100644 index 0000000000000000000000000000000000000000..e711c23780e5e3769b80270bf5b97f6d9f56abbd --- /dev/null +++ b/api/core/tools/provider/builtin/lark_spreadsheet/tools/list_spreadsheet_sheets.py @@ -0,0 +1,18 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class ListSpreadsheetSheetsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + spreadsheet_token = tool_parameters.get("spreadsheet_token") + + res = client.list_spreadsheet_sheets(spreadsheet_token) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_spreadsheet/tools/list_spreadsheet_sheets.yaml b/api/core/tools/provider/builtin/lark_spreadsheet/tools/list_spreadsheet_sheets.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c6a7ef45d46589178ccd4c8d2f145559aa2f5cd0 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_spreadsheet/tools/list_spreadsheet_sheets.yaml @@ -0,0 +1,23 @@ +identity: + name: list_spreadsheet_sheets + author: Doug Lea + label: + en_US: List Spreadsheet Sheets + zh_Hans: 列出电子表格所有工作表 +description: + human: + en_US: List Spreadsheet Sheets + zh_Hans: 列出电子表格所有工作表 + llm: A tool for listing all sheets in a spreadsheet. (列出电子表格所有工作表) +parameters: + - name: spreadsheet_token + type: string + required: true + label: + en_US: Spreadsheet Token + zh_Hans: 电子表格 token + human_description: + en_US: Spreadsheet token, supports input of spreadsheet URL. + zh_Hans: 电子表格 token,支持输入电子表格 URL。 + llm_description: 电子表格 token,支持输入电子表格 URL。 + form: llm diff --git a/api/core/tools/provider/builtin/lark_spreadsheet/tools/read_cols.py b/api/core/tools/provider/builtin/lark_spreadsheet/tools/read_cols.py new file mode 100644 index 0000000000000000000000000000000000000000..1df289c1d71b0104c3d2274a62c60ae854b8ada0 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_spreadsheet/tools/read_cols.py @@ -0,0 +1,23 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class ReadColsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + spreadsheet_token = tool_parameters.get("spreadsheet_token") + sheet_id = tool_parameters.get("sheet_id") + sheet_name = tool_parameters.get("sheet_name") + start_col = tool_parameters.get("start_col") + num_cols = tool_parameters.get("num_cols") + user_id_type = tool_parameters.get("user_id_type", "open_id") + + res = client.read_cols(spreadsheet_token, sheet_id, sheet_name, start_col, num_cols, user_id_type) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_spreadsheet/tools/read_cols.yaml b/api/core/tools/provider/builtin/lark_spreadsheet/tools/read_cols.yaml new file mode 100644 index 0000000000000000000000000000000000000000..34da74592d589864b4144aaba4b4f0777ea660d2 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_spreadsheet/tools/read_cols.yaml @@ -0,0 +1,97 @@ +identity: + name: read_cols + author: Doug Lea + label: + en_US: Read Cols + zh_Hans: 读取工作表列数据 +description: + human: + en_US: Read Cols + zh_Hans: 读取工作表列数据 + llm: A tool for reading column data from a spreadsheet. (读取工作表列数据) +parameters: + - name: spreadsheet_token + type: string + required: true + label: + en_US: spreadsheet_token + zh_Hans: 电子表格 token + human_description: + en_US: Spreadsheet token, supports input of spreadsheet URL. + zh_Hans: 电子表格 token,支持输入电子表格 url。 + llm_description: 电子表格 token,支持输入电子表格 url。 + form: llm + + - name: sheet_id + type: string + required: false + label: + en_US: sheet_id + zh_Hans: 工作表 ID + human_description: + en_US: Sheet ID, either sheet_id or sheet_name must be filled. + zh_Hans: 工作表 ID,与 sheet_name 二者其一必填。 + llm_description: 工作表 ID,与 sheet_name 二者其一必填。 + form: llm + + - name: sheet_name + type: string + required: false + label: + en_US: sheet_name + zh_Hans: 工作表名称 + human_description: + en_US: Sheet name, either sheet_id or sheet_name must be filled. + zh_Hans: 工作表名称,与 sheet_id 二者其一必填。 + llm_description: 工作表名称,与 sheet_id 二者其一必填。 + form: llm + + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form + + - name: start_col + type: number + required: false + label: + en_US: start_col + zh_Hans: 起始列号 + human_description: + en_US: Starting column number, starting from 1. + zh_Hans: 起始列号,从 1 开始。 + llm_description: 起始列号,从 1 开始。 + form: form + + - name: num_cols + type: number + required: true + label: + en_US: num_cols + zh_Hans: 读取列数 + human_description: + en_US: Number of columns to read. + zh_Hans: 读取列数 + llm_description: 读取列数 + form: form diff --git a/api/core/tools/provider/builtin/lark_spreadsheet/tools/read_rows.py b/api/core/tools/provider/builtin/lark_spreadsheet/tools/read_rows.py new file mode 100644 index 0000000000000000000000000000000000000000..1cab38a45452698976ca4e0fbb1fb18e38fc3a66 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_spreadsheet/tools/read_rows.py @@ -0,0 +1,23 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class ReadRowsTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + spreadsheet_token = tool_parameters.get("spreadsheet_token") + sheet_id = tool_parameters.get("sheet_id") + sheet_name = tool_parameters.get("sheet_name") + start_row = tool_parameters.get("start_row") + num_rows = tool_parameters.get("num_rows") + user_id_type = tool_parameters.get("user_id_type", "open_id") + + res = client.read_rows(spreadsheet_token, sheet_id, sheet_name, start_row, num_rows, user_id_type) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_spreadsheet/tools/read_rows.yaml b/api/core/tools/provider/builtin/lark_spreadsheet/tools/read_rows.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5dfa8d5835412561565ab3211634167cab52c39b --- /dev/null +++ b/api/core/tools/provider/builtin/lark_spreadsheet/tools/read_rows.yaml @@ -0,0 +1,97 @@ +identity: + name: read_rows + author: Doug Lea + label: + en_US: Read Rows + zh_Hans: 读取工作表行数据 +description: + human: + en_US: Read Rows + zh_Hans: 读取工作表行数据 + llm: A tool for reading row data from a spreadsheet. (读取工作表行数据) +parameters: + - name: spreadsheet_token + type: string + required: true + label: + en_US: spreadsheet_token + zh_Hans: 电子表格 token + human_description: + en_US: Spreadsheet token, supports input of spreadsheet URL. + zh_Hans: 电子表格 token,支持输入电子表格 url。 + llm_description: 电子表格 token,支持输入电子表格 url。 + form: llm + + - name: sheet_id + type: string + required: false + label: + en_US: sheet_id + zh_Hans: 工作表 ID + human_description: + en_US: Sheet ID, either sheet_id or sheet_name must be filled. + zh_Hans: 工作表 ID,与 sheet_name 二者其一必填。 + llm_description: 工作表 ID,与 sheet_name 二者其一必填。 + form: llm + + - name: sheet_name + type: string + required: false + label: + en_US: sheet_name + zh_Hans: 工作表名称 + human_description: + en_US: Sheet name, either sheet_id or sheet_name must be filled. + zh_Hans: 工作表名称,与 sheet_id 二者其一必填。 + llm_description: 工作表名称,与 sheet_id 二者其一必填。 + form: llm + + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form + + - name: start_row + type: number + required: false + label: + en_US: start_row + zh_Hans: 起始行号 + human_description: + en_US: Starting row number, starting from 1. + zh_Hans: 起始行号,从 1 开始。 + llm_description: 起始行号,从 1 开始。 + form: form + + - name: num_rows + type: number + required: true + label: + en_US: num_rows + zh_Hans: 读取行数 + human_description: + en_US: Number of rows to read. + zh_Hans: 读取行数 + llm_description: 读取行数 + form: form diff --git a/api/core/tools/provider/builtin/lark_spreadsheet/tools/read_table.py b/api/core/tools/provider/builtin/lark_spreadsheet/tools/read_table.py new file mode 100644 index 0000000000000000000000000000000000000000..0f05249004ee207bc3245a0b50ba99ccf24d16a6 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_spreadsheet/tools/read_table.py @@ -0,0 +1,23 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class ReadTableTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + spreadsheet_token = tool_parameters.get("spreadsheet_token") + sheet_id = tool_parameters.get("sheet_id") + sheet_name = tool_parameters.get("sheet_name") + num_range = tool_parameters.get("num_range") + query = tool_parameters.get("query") + user_id_type = tool_parameters.get("user_id_type", "open_id") + + res = client.read_table(spreadsheet_token, sheet_id, sheet_name, num_range, query, user_id_type) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_spreadsheet/tools/read_table.yaml b/api/core/tools/provider/builtin/lark_spreadsheet/tools/read_table.yaml new file mode 100644 index 0000000000000000000000000000000000000000..10534436d66e7a68c63fa191b17db8301ce4661e --- /dev/null +++ b/api/core/tools/provider/builtin/lark_spreadsheet/tools/read_table.yaml @@ -0,0 +1,122 @@ +identity: + name: read_table + author: Doug Lea + label: + en_US: Read Table + zh_Hans: 自定义读取电子表格行列数据 +description: + human: + en_US: Read Table + zh_Hans: 自定义读取电子表格行列数据 + llm: A tool for custom reading of row and column data from a spreadsheet. (自定义读取电子表格行列数据) +parameters: + - name: spreadsheet_token + type: string + required: true + label: + en_US: spreadsheet_token + zh_Hans: 电子表格 token + human_description: + en_US: Spreadsheet token, supports input of spreadsheet URL. + zh_Hans: 电子表格 token,支持输入电子表格 url。 + llm_description: 电子表格 token,支持输入电子表格 url。 + form: llm + + - name: sheet_id + type: string + required: false + label: + en_US: sheet_id + zh_Hans: 工作表 ID + human_description: + en_US: Sheet ID, either sheet_id or sheet_name must be filled. + zh_Hans: 工作表 ID,与 sheet_name 二者其一必填。 + llm_description: 工作表 ID,与 sheet_name 二者其一必填。 + form: llm + + - name: sheet_name + type: string + required: false + label: + en_US: sheet_name + zh_Hans: 工作表名称 + human_description: + en_US: Sheet name, either sheet_id or sheet_name must be filled. + zh_Hans: 工作表名称,与 sheet_id 二者其一必填。 + llm_description: 工作表名称,与 sheet_id 二者其一必填。 + form: llm + + - name: user_id_type + type: select + required: false + options: + - value: open_id + label: + en_US: open_id + zh_Hans: open_id + - value: union_id + label: + en_US: union_id + zh_Hans: union_id + - value: user_id + label: + en_US: user_id + zh_Hans: user_id + default: "open_id" + label: + en_US: user_id_type + zh_Hans: 用户 ID 类型 + human_description: + en_US: User ID type, optional values are open_id, union_id, user_id, with a default value of open_id. + zh_Hans: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + llm_description: 用户 ID 类型,可选值有 open_id、union_id、user_id,默认值为 open_id。 + form: form + + - name: start_row + type: number + required: false + label: + en_US: start_row + zh_Hans: 起始行号 + human_description: + en_US: Starting row number, starting from 1. + zh_Hans: 起始行号,从 1 开始。 + llm_description: 起始行号,从 1 开始。 + form: form + + - name: num_rows + type: number + required: false + label: + en_US: num_rows + zh_Hans: 读取行数 + human_description: + en_US: Number of rows to read. + zh_Hans: 读取行数 + llm_description: 读取行数 + form: form + + - name: range + type: string + required: false + label: + en_US: range + zh_Hans: 取数范围 + human_description: + en_US: | + Data range, format like: A1:B2, can be empty when query=all. + zh_Hans: 取数范围,格式如:A1:B2,query=all 时可为空。 + llm_description: 取数范围,格式如:A1:B2,query=all 时可为空。 + form: llm + + - name: query + type: string + required: false + label: + en_US: query + zh_Hans: 查询 + human_description: + en_US: Pass "all" to query all data in the table, but no more than 100 columns. + zh_Hans: 传 all,表示查询表格所有数据,但最多查询 100 列数据。 + llm_description: 传 all,表示查询表格所有数据,但最多查询 100 列数据。 + form: llm diff --git a/api/core/tools/provider/builtin/lark_task/_assets/icon.png b/api/core/tools/provider/builtin/lark_task/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..26ea6a2eefa5bea073ed8ca8c4d712852745026c Binary files /dev/null and b/api/core/tools/provider/builtin/lark_task/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/lark_task/lark_task.py b/api/core/tools/provider/builtin/lark_task/lark_task.py new file mode 100644 index 0000000000000000000000000000000000000000..02cf009f017e617e054b6ea31304664f7debdb2f --- /dev/null +++ b/api/core/tools/provider/builtin/lark_task/lark_task.py @@ -0,0 +1,7 @@ +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController +from core.tools.utils.lark_api_utils import lark_auth + + +class LarkTaskProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + lark_auth(credentials) diff --git a/api/core/tools/provider/builtin/lark_task/lark_task.yaml b/api/core/tools/provider/builtin/lark_task/lark_task.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ada068b0aab3ce5bfcf5570ae1089260b5202a21 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_task/lark_task.yaml @@ -0,0 +1,36 @@ +identity: + author: Doug Lea + name: lark_task + label: + en_US: Lark Task + zh_Hans: Lark 任务 + description: + en_US: | + Lark Task, requires the following permissions: task:task:write、contact:user.id:readonly. + zh_Hans: | + Lark 任务,需要开通以下权限: task:task:write、contact:user.id:readonly。 + icon: icon.png + tags: + - social + - productivity +credentials_for_provider: + app_id: + type: text-input + required: true + label: + en_US: APP ID + placeholder: + en_US: Please input your Lark app id + zh_Hans: 请输入你的 Lark app id + help: + en_US: Get your app_id and app_secret from Lark + zh_Hans: 从 Lark 获取您的 app_id 和 app_secret + url: https://open.larksuite.com/app + app_secret: + type: secret-input + required: true + label: + en_US: APP Secret + placeholder: + en_US: Please input your app secret + zh_Hans: 请输入你的 Lark app secret diff --git a/api/core/tools/provider/builtin/lark_task/tools/add_members.py b/api/core/tools/provider/builtin/lark_task/tools/add_members.py new file mode 100644 index 0000000000000000000000000000000000000000..9b8e4d68f394a822717460054705fb20cae81cd6 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_task/tools/add_members.py @@ -0,0 +1,20 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class AddMembersTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + task_guid = tool_parameters.get("task_guid") + member_phone_or_email = tool_parameters.get("member_phone_or_email") + member_role = tool_parameters.get("member_role", "follower") + + res = client.add_members(task_guid, member_phone_or_email, member_role) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_task/tools/add_members.yaml b/api/core/tools/provider/builtin/lark_task/tools/add_members.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0b12172e0b85e79aa062eafe1624f129bf0bd601 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_task/tools/add_members.yaml @@ -0,0 +1,58 @@ +identity: + name: add_members + author: Doug Lea + label: + en_US: Add Lark Members + zh_Hans: 添加 Lark 任务成员 +description: + human: + en_US: Add Lark Members + zh_Hans: 添加 Lark 任务成员 + llm: A tool for adding members to a Lark task.(添加 Lark 任务成员) +parameters: + - name: task_guid + type: string + required: true + label: + en_US: Task GUID + zh_Hans: 任务 GUID + human_description: + en_US: | + The GUID of the task to be added, supports passing either the Task ID or the Task link URL. Example of Task ID: 8b5425ec-9f2a-43bd-a3ab-01912f50282b; Example of Task link URL: https://applink.larksuite.com/client/todo/detail?guid=1b066afa-96de-406c-90a3-dfd30159a571&suite_entity_num=t100805 + zh_Hans: 要添加的任务的 GUID,支持传任务 ID 和任务链接 URL。任务 ID 示例:8b5425ec-9f2a-43bd-a3ab-01912f50282b;任务链接 URL 示例:https://applink.larksuite.com/client/todo/detail?guid=1b066afa-96de-406c-90a3-dfd30159a571&suite_entity_num=t100805 + llm_description: 要添加的任务的 GUID,支持传任务 ID 和任务链接 URL。任务 ID 示例:8b5425ec-9f2a-43bd-a3ab-01912f50282b;任务链接 URL 示例:https://applink.larksuite.com/client/todo/detail?guid=1b066afa-96de-406c-90a3-dfd30159a571&suite_entity_num=t100805 + form: llm + + - name: member_phone_or_email + type: string + required: true + label: + en_US: Task Member Phone Or Email + zh_Hans: 任务成员的电话或邮箱 + human_description: + en_US: A list of member emails or phone numbers, separated by commas. + zh_Hans: 任务成员邮箱或者手机号列表,使用逗号分隔。 + llm_description: 任务成员邮箱或者手机号列表,使用逗号分隔。 + form: llm + + - name: member_role + type: select + required: true + options: + - value: assignee + label: + en_US: assignee + zh_Hans: 负责人 + - value: follower + label: + en_US: follower + zh_Hans: 关注人 + default: "follower" + label: + en_US: member_role + zh_Hans: 成员的角色 + human_description: + en_US: Member role, optional values are "assignee" (responsible person) and "follower" (observer), with a default value of "assignee". + zh_Hans: 成员的角色,可选值有 "assignee"(负责人)和 "follower"(关注人),默认值为 "assignee"。 + llm_description: 成员的角色,可选值有 "assignee"(负责人)和 "follower"(关注人),默认值为 "assignee"。 + form: form diff --git a/api/core/tools/provider/builtin/lark_task/tools/create_task.py b/api/core/tools/provider/builtin/lark_task/tools/create_task.py new file mode 100644 index 0000000000000000000000000000000000000000..ff37593fbe3a12b1ce5749ebcdd0db7fa1e79f28 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_task/tools/create_task.py @@ -0,0 +1,22 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class CreateTaskTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + summary = tool_parameters.get("summary") + start_time = tool_parameters.get("start_time") + end_time = tool_parameters.get("end_time") + completed_time = tool_parameters.get("completed_time") + description = tool_parameters.get("description") + + res = client.create_task(summary, start_time, end_time, completed_time, description) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_task/tools/create_task.yaml b/api/core/tools/provider/builtin/lark_task/tools/create_task.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4303763a1dd40631d26ed70deee80c1159631745 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_task/tools/create_task.yaml @@ -0,0 +1,74 @@ +identity: + name: create_task + author: Doug Lea + label: + en_US: Create Lark Task + zh_Hans: 创建 Lark 任务 +description: + human: + en_US: Create Lark Task + zh_Hans: 创建 Lark 任务 + llm: A tool for creating tasks in Lark.(创建 Lark 任务) +parameters: + - name: summary + type: string + required: true + label: + en_US: Task Title + zh_Hans: 任务标题 + human_description: + en_US: The title of the task. + zh_Hans: 任务标题 + llm_description: 任务标题 + form: llm + + - name: description + type: string + required: false + label: + en_US: Task Description + zh_Hans: 任务备注 + human_description: + en_US: The description or notes for the task. + zh_Hans: 任务备注 + llm_description: 任务备注 + form: llm + + - name: start_time + type: string + required: false + label: + en_US: Start Time + zh_Hans: 任务开始时间 + human_description: + en_US: | + The start time of the task, in the format: 2006-01-02 15:04:05 + zh_Hans: 任务开始时间,格式为:2006-01-02 15:04:05 + llm_description: 任务开始时间,格式为:2006-01-02 15:04:05 + form: llm + + - name: end_time + type: string + required: false + label: + en_US: End Time + zh_Hans: 任务结束时间 + human_description: + en_US: | + The end time of the task, in the format: 2006-01-02 15:04:05 + zh_Hans: 任务结束时间,格式为:2006-01-02 15:04:05 + llm_description: 任务结束时间,格式为:2006-01-02 15:04:05 + form: llm + + - name: completed_time + type: string + required: false + label: + en_US: Completed Time + zh_Hans: 任务完成时间 + human_description: + en_US: | + The completion time of the task, in the format: 2006-01-02 15:04:05. Leave empty to create an incomplete task; fill in a specific time to create a completed task. + zh_Hans: 任务完成时间,格式为:2006-01-02 15:04:05,不填写表示创建一个未完成任务;填写一个具体的时间表示创建一个已完成任务。 + llm_description: 任务完成时间,格式为:2006-01-02 15:04:05,不填写表示创建一个未完成任务;填写一个具体的时间表示创建一个已完成任务。 + form: llm diff --git a/api/core/tools/provider/builtin/lark_task/tools/delete_task.py b/api/core/tools/provider/builtin/lark_task/tools/delete_task.py new file mode 100644 index 0000000000000000000000000000000000000000..eca381be2c185e70b2c36b528ebbb4df1811fd48 --- /dev/null +++ b/api/core/tools/provider/builtin/lark_task/tools/delete_task.py @@ -0,0 +1,18 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class UpdateTaskTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + task_guid = tool_parameters.get("task_guid") + + res = client.delete_task(task_guid) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_task/tools/delete_task.yaml b/api/core/tools/provider/builtin/lark_task/tools/delete_task.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bc0154d9dc5c774560e48ddc95ac43dfb00445ed --- /dev/null +++ b/api/core/tools/provider/builtin/lark_task/tools/delete_task.yaml @@ -0,0 +1,24 @@ +identity: + name: delete_task + author: Doug Lea + label: + en_US: Delete Lark Task + zh_Hans: 删除 Lark 任务 +description: + human: + en_US: Delete Lark Task + zh_Hans: 删除 Lark 任务 + llm: A tool for deleting tasks in Lark.(删除 Lark 任务) +parameters: + - name: task_guid + type: string + required: true + label: + en_US: Task GUID + zh_Hans: 任务 GUID + human_description: + en_US: | + The GUID of the task to be deleted, supports passing either the Task ID or the Task link URL. Example of Task ID: 8b5425ec-9f2a-43bd-a3ab-01912f50282b; Example of Task link URL: https://applink.larksuite.com/client/todo/detail?guid=1b066afa-96de-406c-90a3-dfd30159a571&suite_entity_num=t100805 + zh_Hans: 要删除的任务的 GUID,支持传任务 ID 和任务链接 URL。任务 ID 示例:8b5425ec-9f2a-43bd-a3ab-01912f50282b;任务链接 URL 示例:https://applink.larksuite.com/client/todo/detail?guid=1b066afa-96de-406c-90a3-dfd30159a571&suite_entity_num=t100805 + llm_description: 要删除的任务的 GUID,支持传任务 ID 和任务链接 URL。任务 ID 示例:8b5425ec-9f2a-43bd-a3ab-01912f50282b;任务链接 URL 示例:https://applink.larksuite.com/client/todo/detail?guid=1b066afa-96de-406c-90a3-dfd30159a571&suite_entity_num=t100805 + form: llm diff --git a/api/core/tools/provider/builtin/lark_task/tools/update_task.py b/api/core/tools/provider/builtin/lark_task/tools/update_task.py new file mode 100644 index 0000000000000000000000000000000000000000..0d3469c91a01bc17ca88f0807ca4dea6f1985c6d --- /dev/null +++ b/api/core/tools/provider/builtin/lark_task/tools/update_task.py @@ -0,0 +1,23 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class UpdateTaskTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + task_guid = tool_parameters.get("task_guid") + summary = tool_parameters.get("summary") + start_time = tool_parameters.get("start_time") + end_time = tool_parameters.get("end_time") + completed_time = tool_parameters.get("completed_time") + description = tool_parameters.get("description") + + res = client.update_task(task_guid, summary, start_time, end_time, completed_time, description) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_task/tools/update_task.yaml b/api/core/tools/provider/builtin/lark_task/tools/update_task.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a98f037f211c9a4f01e5c7b76949265df71154ca --- /dev/null +++ b/api/core/tools/provider/builtin/lark_task/tools/update_task.yaml @@ -0,0 +1,89 @@ +identity: + name: update_task + author: Doug Lea + label: + en_US: Update Lark Task + zh_Hans: 更新 Lark 任务 +description: + human: + en_US: Update Lark Task + zh_Hans: 更新 Lark 任务 + llm: A tool for updating tasks in Lark.(更新 Lark 任务) +parameters: + - name: task_guid + type: string + required: true + label: + en_US: Task GUID + zh_Hans: 任务 GUID + human_description: + en_US: | + The task ID, supports inputting either the Task ID or the Task link URL. Example of Task ID: 42cad8a0-f8c8-4344-9be2-d1d7e8e91b64; Example of Task link URL: https://applink.larksuite.com/client/todo/detail?guid=1b066afa-96de-406c-90a3-dfd30159a571&suite_entity_num=t100805 + zh_Hans: | + 任务ID,支持传入任务 ID 和任务链接 URL。任务 ID 示例: 42cad8a0-f8c8-4344-9be2-d1d7e8e91b64;任务链接 URL 示例: https://applink.larksuite.com/client/todo/detail?guid=1b066afa-96de-406c-90a3-dfd30159a571&suite_entity_num=t100805 + llm_description: | + 任务ID,支持传入任务 ID 和任务链接 URL。任务 ID 示例: 42cad8a0-f8c8-4344-9be2-d1d7e8e91b64;任务链接 URL 示例: https://applink.larksuite.com/client/todo/detail?guid=1b066afa-96de-406c-90a3-dfd30159a571&suite_entity_num=t100805 + form: llm + + - name: summary + type: string + required: true + label: + en_US: Task Title + zh_Hans: 任务标题 + human_description: + en_US: The title of the task. + zh_Hans: 任务标题 + llm_description: 任务标题 + form: llm + + - name: description + type: string + required: false + label: + en_US: Task Description + zh_Hans: 任务备注 + human_description: + en_US: The description or notes for the task. + zh_Hans: 任务备注 + llm_description: 任务备注 + form: llm + + - name: start_time + type: string + required: false + label: + en_US: Start Time + zh_Hans: 任务开始时间 + human_description: + en_US: | + The start time of the task, in the format: 2006-01-02 15:04:05 + zh_Hans: 任务开始时间,格式为:2006-01-02 15:04:05 + llm_description: 任务开始时间,格式为:2006-01-02 15:04:05 + form: llm + + - name: end_time + type: string + required: false + label: + en_US: End Time + zh_Hans: 任务结束时间 + human_description: + en_US: | + The end time of the task, in the format: 2006-01-02 15:04:05 + zh_Hans: 任务结束时间,格式为:2006-01-02 15:04:05 + llm_description: 任务结束时间,格式为:2006-01-02 15:04:05 + form: llm + + - name: completed_time + type: string + required: false + label: + en_US: Completed Time + zh_Hans: 任务完成时间 + human_description: + en_US: | + The completion time of the task, in the format: 2006-01-02 15:04:05 + zh_Hans: 任务完成时间,格式为:2006-01-02 15:04:05 + llm_description: 任务完成时间,格式为:2006-01-02 15:04:05 + form: llm diff --git a/api/core/tools/provider/builtin/lark_wiki/_assets/icon.png b/api/core/tools/provider/builtin/lark_wiki/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..47f6b8c30ea0cf2d4b625bf5890ab94b0b5ca84e Binary files /dev/null and b/api/core/tools/provider/builtin/lark_wiki/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/lark_wiki/lark_wiki.py b/api/core/tools/provider/builtin/lark_wiki/lark_wiki.py new file mode 100644 index 0000000000000000000000000000000000000000..e6941206ee761823ef469a3caa29f8765e0a96ea --- /dev/null +++ b/api/core/tools/provider/builtin/lark_wiki/lark_wiki.py @@ -0,0 +1,7 @@ +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController +from core.tools.utils.lark_api_utils import lark_auth + + +class LarkWikiProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + lark_auth(credentials) diff --git a/api/core/tools/provider/builtin/lark_wiki/lark_wiki.yaml b/api/core/tools/provider/builtin/lark_wiki/lark_wiki.yaml new file mode 100644 index 0000000000000000000000000000000000000000..86bef000868d680cdc8a1b46bd174fccd8c00f9a --- /dev/null +++ b/api/core/tools/provider/builtin/lark_wiki/lark_wiki.yaml @@ -0,0 +1,36 @@ +identity: + author: Doug Lea + name: lark_wiki + label: + en_US: Lark Wiki + zh_Hans: Lark 知识库 + description: + en_US: | + Lark Wiki, requires the following permissions: wiki:wiki:readonly. + zh_Hans: | + Lark 知识库,需要开通以下权限: wiki:wiki:readonly。 + icon: icon.png + tags: + - social + - productivity +credentials_for_provider: + app_id: + type: text-input + required: true + label: + en_US: APP ID + placeholder: + en_US: Please input your Lark app id + zh_Hans: 请输入你的 Lark app id + help: + en_US: Get your app_id and app_secret from Lark + zh_Hans: 从 Lark 获取您的 app_id 和 app_secret + url: https://open.larksuite.com/app + app_secret: + type: secret-input + required: true + label: + en_US: APP Secret + placeholder: + en_US: Please input your app secret + zh_Hans: 请输入你的 Lark app secret diff --git a/api/core/tools/provider/builtin/lark_wiki/tools/get_wiki_nodes.py b/api/core/tools/provider/builtin/lark_wiki/tools/get_wiki_nodes.py new file mode 100644 index 0000000000000000000000000000000000000000..a05f300755962f07dd6cc9dbb258d953df40bdda --- /dev/null +++ b/api/core/tools/provider/builtin/lark_wiki/tools/get_wiki_nodes.py @@ -0,0 +1,21 @@ +from typing import Any + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.lark_api_utils import LarkRequest + + +class GetWikiNodesTool(BuiltinTool): + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage: + app_id = self.runtime.credentials.get("app_id") + app_secret = self.runtime.credentials.get("app_secret") + client = LarkRequest(app_id, app_secret) + + space_id = tool_parameters.get("space_id") + parent_node_token = tool_parameters.get("parent_node_token") + page_token = tool_parameters.get("page_token") + page_size = tool_parameters.get("page_size") + + res = client.get_wiki_nodes(space_id, parent_node_token, page_token, page_size) + + return self.create_json_message(res) diff --git a/api/core/tools/provider/builtin/lark_wiki/tools/get_wiki_nodes.yaml b/api/core/tools/provider/builtin/lark_wiki/tools/get_wiki_nodes.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a8c242a2e909df3f70ee0f148ae573f070eeeaea --- /dev/null +++ b/api/core/tools/provider/builtin/lark_wiki/tools/get_wiki_nodes.yaml @@ -0,0 +1,63 @@ +identity: + name: get_wiki_nodes + author: Doug Lea + label: + en_US: Get Wiki Nodes + zh_Hans: 获取知识空间子节点列表 +description: + human: + en_US: | + Get the list of child nodes in Wiki, make sure the app/bot is a member of the wiki space. See How to add an app as a wiki base administrator (member). https://open.larksuite.com/document/server-docs/docs/wiki-v2/wiki-qa + zh_Hans: | + 获取知识库全部子节点列表,请确保应用/机器人为知识空间成员。参阅如何将应用添加为知识库管理员(成员)。https://open.larksuite.com/document/server-docs/docs/wiki-v2/wiki-qa + llm: A tool for getting all sub-nodes of a knowledge base.(获取知识空间子节点列表) +parameters: + - name: space_id + type: string + required: true + label: + en_US: Space Id + zh_Hans: 知识空间 ID + human_description: + en_US: | + The ID of the knowledge space. Supports space link URL, for example: https://lark-japan.jp.larksuite.com/wiki/settings/7431084851517718561 + zh_Hans: 知识空间 ID,支持空间链接 URL,例如:https://lark-japan.jp.larksuite.com/wiki/settings/7431084851517718561 + llm_description: 知识空间 ID,支持空间链接 URL,例如:https://lark-japan.jp.larksuite.com/wiki/settings/7431084851517718561 + form: llm + + - name: page_size + type: number + required: false + default: 10 + label: + en_US: Page Size + zh_Hans: 分页大小 + human_description: + en_US: The size of each page, with a maximum value of 50. + zh_Hans: 分页大小,最大值 50。 + llm_description: 分页大小,最大值 50。 + form: form + + - name: page_token + type: string + required: false + label: + en_US: Page Token + zh_Hans: 分页标记 + human_description: + en_US: The pagination token. Leave empty for the first request to start from the beginning; if the paginated query result has more items, a new page_token will be returned, which can be used to get the next set of results. + zh_Hans: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + llm_description: 分页标记,第一次请求不填,表示从头开始遍历;分页查询结果还有更多项时会同时返回新的 page_token,下次遍历可采用该 page_token 获取查询结果。 + form: llm + + - name: parent_node_token + type: string + required: false + label: + en_US: Parent Node Token + zh_Hans: 父节点 token + human_description: + en_US: The token of the parent node. + zh_Hans: 父节点 token + llm_description: 父节点 token + form: llm diff --git a/api/core/tools/provider/builtin/maths/_assets/icon.svg b/api/core/tools/provider/builtin/maths/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..f94d1152113830d08d7d55b126644db0360a7882 --- /dev/null +++ b/api/core/tools/provider/builtin/maths/_assets/icon.svg @@ -0,0 +1,10 @@ + + + + + + + + + + diff --git a/api/core/tools/provider/builtin/maths/maths.py b/api/core/tools/provider/builtin/maths/maths.py new file mode 100644 index 0000000000000000000000000000000000000000..d4b449ec87a18ae170ff8ebd09c869c45d7c92e8 --- /dev/null +++ b/api/core/tools/provider/builtin/maths/maths.py @@ -0,0 +1,18 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.maths.tools.eval_expression import EvaluateExpressionTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class MathsProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + EvaluateExpressionTool().invoke( + user_id="", + tool_parameters={ + "expression": "1+(2+3)*4", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/maths/maths.yaml b/api/core/tools/provider/builtin/maths/maths.yaml new file mode 100644 index 0000000000000000000000000000000000000000..35c2380e29a701af1323ad20e1c11be3d2a923dd --- /dev/null +++ b/api/core/tools/provider/builtin/maths/maths.yaml @@ -0,0 +1,15 @@ +identity: + author: Bowen Liang + name: maths + label: + en_US: Maths + zh_Hans: 数学工具 + pt_BR: Maths + description: + en_US: A tool for maths. + zh_Hans: 一个用于数学计算的工具。 + pt_BR: A tool for maths. + icon: icon.svg + tags: + - utilities + - productivity diff --git a/api/core/tools/provider/builtin/maths/tools/eval_expression.py b/api/core/tools/provider/builtin/maths/tools/eval_expression.py new file mode 100644 index 0000000000000000000000000000000000000000..d3a497d1cd5c54ebd5cfdcd4baa541e45cfa5997 --- /dev/null +++ b/api/core/tools/provider/builtin/maths/tools/eval_expression.py @@ -0,0 +1,30 @@ +import logging +from typing import Any, Union + +import numexpr as ne # type: ignore + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class EvaluateExpressionTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + # get expression + expression = tool_parameters.get("expression", "").strip() + if not expression: + return self.create_text_message("Invalid expression") + + try: + result = ne.evaluate(expression) + result_str = str(result) + except Exception as e: + logging.exception(f"Error evaluating expression: {expression}") + return self.create_text_message(f"Invalid expression: {expression}, error: {str(e)}") + return self.create_text_message(f'The result of the expression "{expression}" is {result_str}') diff --git a/api/core/tools/provider/builtin/maths/tools/eval_expression.yaml b/api/core/tools/provider/builtin/maths/tools/eval_expression.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c936a4293fbe72aae1aa7620518180b86f85b39d --- /dev/null +++ b/api/core/tools/provider/builtin/maths/tools/eval_expression.yaml @@ -0,0 +1,26 @@ +identity: + name: eval_expression + author: Bowen Liang + label: + en_US: Evaluate Math Expression + zh_Hans: 计算数学表达式 + pt_BR: Evaluate Math Expression +description: + human: + en_US: A tool for evaluating an math expression, calculated locally with NumExpr. + zh_Hans: 一个用于计算数学表达式的工具,表达式将通过NumExpr本地执行。 + pt_BR: A tool for evaluating an math expression, calculated locally with NumExpr. + llm: A tool for evaluating an math expression. +parameters: + - name: expression + type: string + required: true + label: + en_US: Math Expression + zh_Hans: 数学计算表达式 + pt_BR: Math Expression + human_description: + en_US: Math Expression + zh_Hans: 数学计算表达式 + pt_BR: Math Expression + form: llm diff --git a/api/core/tools/provider/builtin/nominatim/_assets/icon.svg b/api/core/tools/provider/builtin/nominatim/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..db5a4eb868c5e8c14b11d473ec4c06817a962974 --- /dev/null +++ b/api/core/tools/provider/builtin/nominatim/_assets/icon.svg @@ -0,0 +1,277 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + 010110010011010110010011 + 010110010011010110010011 + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/nominatim/nominatim.py b/api/core/tools/provider/builtin/nominatim/nominatim.py new file mode 100644 index 0000000000000000000000000000000000000000..5a24bed7507eb64812d12ddbdc29facf50b4fd1a --- /dev/null +++ b/api/core/tools/provider/builtin/nominatim/nominatim.py @@ -0,0 +1,27 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.nominatim.tools.nominatim_search import NominatimSearchTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class NominatimProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + result = ( + NominatimSearchTool() + .fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ) + .invoke( + user_id="", + tool_parameters={ + "query": "London", + "limit": 1, + }, + ) + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/nominatim/nominatim.yaml b/api/core/tools/provider/builtin/nominatim/nominatim.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7d014bd78c6a59f0eeef2a699bf4145b4141ac93 --- /dev/null +++ b/api/core/tools/provider/builtin/nominatim/nominatim.yaml @@ -0,0 +1,43 @@ +identity: + author: Charles Zhou + name: nominatim + label: + en_US: Nominatim + zh_Hans: Nominatim + de_DE: Nominatim + ja_JP: Nominatim + description: + en_US: Nominatim is a search engine for OpenStreetMap data + zh_Hans: Nominatim是OpenStreetMap数据的搜索引擎 + de_DE: Nominatim ist eine Suchmaschine für OpenStreetMap-Daten + ja_JP: NominatimはOpenStreetMapデータの検索エンジンです + icon: icon.svg + tags: + - search + - utilities +credentials_for_provider: + base_url: + type: text-input + required: false + default: https://nominatim.openstreetmap.org + label: + en_US: Nominatim Base URL + zh_Hans: Nominatim 基础 URL + de_DE: Nominatim Basis-URL + ja_JP: Nominatim ベースURL + placeholder: + en_US: "Enter your Nominatim instance URL (default: + https://nominatim.openstreetmap.org)" + zh_Hans: 输入您的Nominatim实例URL(默认:https://nominatim.openstreetmap.org) + de_DE: "Geben Sie Ihre Nominatim-Instanz-URL ein (Standard: + https://nominatim.openstreetmap.org)" + ja_JP: NominatimインスタンスのURLを入力してください(デフォルト:https://nominatim.openstreetmap.org) + help: + en_US: The base URL for the Nominatim instance. Use the default for the public + service or enter your self-hosted instance URL. + zh_Hans: Nominatim实例的基础URL。使用默认值可访问公共服务,或输入您的自托管实例URL。 + de_DE: Die Basis-URL für die Nominatim-Instanz. Verwenden Sie den Standardwert + für den öffentlichen Dienst oder geben Sie die URL Ihrer selbst + gehosteten Instanz ein. + ja_JP: NominatimインスタンスのベースURL。公共サービスにはデフォルトを使用するか、自己ホスティングインスタンスのURLを入力してください。 + url: https://nominatim.org/ diff --git a/api/core/tools/provider/builtin/nominatim/tools/nominatim_lookup.py b/api/core/tools/provider/builtin/nominatim/tools/nominatim_lookup.py new file mode 100644 index 0000000000000000000000000000000000000000..ffa8ad0fcc02e09853e3eebfa466778578786d19 --- /dev/null +++ b/api/core/tools/provider/builtin/nominatim/tools/nominatim_lookup.py @@ -0,0 +1,40 @@ +import json +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class NominatimLookupTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + osm_ids = tool_parameters.get("osm_ids", "") + + if not osm_ids: + return self.create_text_message("Please provide OSM IDs") + + params = {"osm_ids": osm_ids, "format": "json", "addressdetails": 1} + + return self._make_request(user_id, "lookup", params) + + def _make_request(self, user_id: str, endpoint: str, params: dict) -> ToolInvokeMessage: + base_url = self.runtime.credentials.get("base_url", "https://nominatim.openstreetmap.org") + + try: + headers = {"User-Agent": "DifyNominatimTool/1.0"} + s = requests.session() + response = s.request(method="GET", headers=headers, url=f"{base_url}/{endpoint}", params=params) + response_data = response.json() + + if response.status_code == 200: + s.close() + return self.create_text_message( + self.summary(user_id=user_id, content=json.dumps(response_data, ensure_ascii=False)) + ) + else: + return self.create_text_message(f"Error: {response.status_code} - {response.text}") + except Exception as e: + return self.create_text_message(f"An error occurred: {str(e)}") diff --git a/api/core/tools/provider/builtin/nominatim/tools/nominatim_lookup.yaml b/api/core/tools/provider/builtin/nominatim/tools/nominatim_lookup.yaml new file mode 100644 index 0000000000000000000000000000000000000000..508c4dcd88ff15c800e5cc9b6e89967fff2e4d07 --- /dev/null +++ b/api/core/tools/provider/builtin/nominatim/tools/nominatim_lookup.yaml @@ -0,0 +1,31 @@ +identity: + name: nominatim_lookup + author: Charles Zhou + label: + en_US: Nominatim OSM Lookup + zh_Hans: Nominatim OSM 对象查找 + de_DE: Nominatim OSM-Objektsuche + ja_JP: Nominatim OSM ルックアップ +description: + human: + en_US: Look up OSM objects using their IDs with Nominatim + zh_Hans: 使用Nominatim通过ID查找OSM对象 + de_DE: Suchen Sie OSM-Objekte anhand ihrer IDs mit Nominatim + ja_JP: Nominatimを使用してIDでOSMオブジェクトを検索 + llm: A tool for looking up OpenStreetMap objects using their IDs with Nominatim. +parameters: + - name: osm_ids + type: string + required: true + label: + en_US: OSM IDs + zh_Hans: OSM ID + de_DE: OSM-IDs + ja_JP: OSM ID + human_description: + en_US: Comma-separated list of OSM IDs to lookup (e.g., N123,W456,R789) + zh_Hans: 要查找的OSM ID的逗号分隔列表(例如:N123,W456,R789) + de_DE: Kommagetrennte Liste von OSM-IDs für die Suche (z.B. N123,W456,R789) + ja_JP: 検索するOSM IDのカンマ区切りリスト(例:N123,W456,R789) + llm_description: A comma-separated list of OSM IDs (prefixed with N, W, or R) for lookup. + form: llm diff --git a/api/core/tools/provider/builtin/nominatim/tools/nominatim_reverse.py b/api/core/tools/provider/builtin/nominatim/tools/nominatim_reverse.py new file mode 100644 index 0000000000000000000000000000000000000000..f46691e1a3ebb4ad896b856f7890f1794c051a1b --- /dev/null +++ b/api/core/tools/provider/builtin/nominatim/tools/nominatim_reverse.py @@ -0,0 +1,41 @@ +import json +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class NominatimReverseTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + lat = tool_parameters.get("lat") + lon = tool_parameters.get("lon") + + if lat is None or lon is None: + return self.create_text_message("Please provide both latitude and longitude") + + params = {"lat": lat, "lon": lon, "format": "json", "addressdetails": 1} + + return self._make_request(user_id, "reverse", params) + + def _make_request(self, user_id: str, endpoint: str, params: dict) -> ToolInvokeMessage: + base_url = self.runtime.credentials.get("base_url", "https://nominatim.openstreetmap.org") + + try: + headers = {"User-Agent": "DifyNominatimTool/1.0"} + s = requests.session() + response = s.request(method="GET", headers=headers, url=f"{base_url}/{endpoint}", params=params) + response_data = response.json() + + if response.status_code == 200: + s.close() + return self.create_text_message( + self.summary(user_id=user_id, content=json.dumps(response_data, ensure_ascii=False)) + ) + else: + return self.create_text_message(f"Error: {response.status_code} - {response.text}") + except Exception as e: + return self.create_text_message(f"An error occurred: {str(e)}") diff --git a/api/core/tools/provider/builtin/nominatim/tools/nominatim_reverse.yaml b/api/core/tools/provider/builtin/nominatim/tools/nominatim_reverse.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f1a2dd09fbc5d55943335e2e915444a92d97d2fa --- /dev/null +++ b/api/core/tools/provider/builtin/nominatim/tools/nominatim_reverse.yaml @@ -0,0 +1,47 @@ +identity: + name: nominatim_reverse + author: Charles Zhou + label: + en_US: Nominatim Reverse Geocoding + zh_Hans: Nominatim 反向地理编码 + de_DE: Nominatim Rückwärts-Geocodierung + ja_JP: Nominatim リバースジオコーディング +description: + human: + en_US: Convert coordinates to addresses using Nominatim + zh_Hans: 使用Nominatim将坐标转换为地址 + de_DE: Konvertieren Sie Koordinaten in Adressen mit Nominatim + ja_JP: Nominatimを使用して座標を住所に変換 + llm: A tool for reverse geocoding using Nominatim, which can convert latitude + and longitude coordinates to an address. +parameters: + - name: lat + type: number + required: true + label: + en_US: Latitude + zh_Hans: 纬度 + de_DE: Breitengrad + ja_JP: 緯度 + human_description: + en_US: Latitude coordinate for reverse geocoding + zh_Hans: 用于反向地理编码的纬度坐标 + de_DE: Breitengrad-Koordinate für die Rückwärts-Geocodierung + ja_JP: リバースジオコーディングの緯度座標 + llm_description: The latitude coordinate for reverse geocoding. + form: llm + - name: lon + type: number + required: true + label: + en_US: Longitude + zh_Hans: 经度 + de_DE: Längengrad + ja_JP: 経度 + human_description: + en_US: Longitude coordinate for reverse geocoding + zh_Hans: 用于反向地理编码的经度坐标 + de_DE: Längengrad-Koordinate für die Rückwärts-Geocodierung + ja_JP: リバースジオコーディングの経度座標 + llm_description: The longitude coordinate for reverse geocoding. + form: llm diff --git a/api/core/tools/provider/builtin/nominatim/tools/nominatim_search.py b/api/core/tools/provider/builtin/nominatim/tools/nominatim_search.py new file mode 100644 index 0000000000000000000000000000000000000000..34851d86dcaa5f38dbaf0f7661d3c797e3f4ac0f --- /dev/null +++ b/api/core/tools/provider/builtin/nominatim/tools/nominatim_search.py @@ -0,0 +1,41 @@ +import json +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class NominatimSearchTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + query = tool_parameters.get("query", "") + limit = tool_parameters.get("limit", 10) + + if not query: + return self.create_text_message("Please input a search query") + + params = {"q": query, "format": "json", "limit": limit, "addressdetails": 1} + + return self._make_request(user_id, "search", params) + + def _make_request(self, user_id: str, endpoint: str, params: dict) -> ToolInvokeMessage: + base_url = self.runtime.credentials.get("base_url", "https://nominatim.openstreetmap.org") + + try: + headers = {"User-Agent": "DifyNominatimTool/1.0"} + s = requests.session() + response = s.request(method="GET", headers=headers, url=f"{base_url}/{endpoint}", params=params) + response_data = response.json() + + if response.status_code == 200: + s.close() + return self.create_text_message( + self.summary(user_id=user_id, content=json.dumps(response_data, ensure_ascii=False)) + ) + else: + return self.create_text_message(f"Error: {response.status_code} - {response.text}") + except Exception as e: + return self.create_text_message(f"An error occurred: {str(e)}") diff --git a/api/core/tools/provider/builtin/nominatim/tools/nominatim_search.yaml b/api/core/tools/provider/builtin/nominatim/tools/nominatim_search.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e0c53c046a2a41e90d1ed2abc64eb3810ef4c9ad --- /dev/null +++ b/api/core/tools/provider/builtin/nominatim/tools/nominatim_search.yaml @@ -0,0 +1,51 @@ +identity: + name: nominatim_search + author: Charles Zhou + label: + en_US: Nominatim Search + zh_Hans: Nominatim 搜索 + de_DE: Nominatim Suche + ja_JP: Nominatim 検索 +description: + human: + en_US: Search for locations using Nominatim + zh_Hans: 使用Nominatim搜索位置 + de_DE: Suche nach Orten mit Nominatim + ja_JP: Nominatimを使用して場所を検索 + llm: A tool for geocoding using Nominatim, which can search for locations based + on addresses or place names. +parameters: + - name: query + type: string + required: true + label: + en_US: Search Query + zh_Hans: 搜索查询 + de_DE: Suchanfrage + ja_JP: 検索クエリ + human_description: + en_US: Enter an address or place name to search for + zh_Hans: 输入要搜索的地址或地名 + de_DE: Geben Sie eine Adresse oder einen Ortsnamen für die Suche ein + ja_JP: 検索する住所または場所の名前を入力してください + llm_description: The search query for Nominatim, which can be an address or place name. + form: llm + - name: limit + type: number + default: 10 + min: 1 + max: 40 + required: false + label: + en_US: Result Limit + zh_Hans: 结果限制 + de_DE: Ergebnislimit + ja_JP: 結果の制限 + human_description: + en_US: "Maximum number of results to return (default: 10, max: 40)" + zh_Hans: 要返回的最大结果数(默认:10,最大:40) + de_DE: "Maximale Anzahl der zurückzugebenden Ergebnisse (Standard: 10, max: 40)" + ja_JP: 返す結果の最大数(デフォルト:10、最大:40) + llm_description: Limit the number of returned results. The default is 10, and + the maximum is 40. + form: form diff --git a/api/core/tools/provider/builtin/novitaai/_assets/icon.ico b/api/core/tools/provider/builtin/novitaai/_assets/icon.ico new file mode 100644 index 0000000000000000000000000000000000000000..e353ecf711cac1d0b1843a9d82793a70a209728f Binary files /dev/null and b/api/core/tools/provider/builtin/novitaai/_assets/icon.ico differ diff --git a/api/core/tools/provider/builtin/novitaai/_novita_tool_base.py b/api/core/tools/provider/builtin/novitaai/_novita_tool_base.py new file mode 100644 index 0000000000000000000000000000000000000000..6473c509e1f4c229e51dea8d9cc14731c304980c --- /dev/null +++ b/api/core/tools/provider/builtin/novitaai/_novita_tool_base.py @@ -0,0 +1,69 @@ +from novita_client import ( # type: ignore + Txt2ImgV3Embedding, + Txt2ImgV3HiresFix, + Txt2ImgV3LoRA, + Txt2ImgV3Refiner, + V3TaskImage, +) + + +class NovitaAiToolBase: + def _extract_loras(self, loras_str: str): + if not loras_str: + return [] + + loras_ori_list = loras_str.strip().split(";") + result_list = [] + for lora_str in loras_ori_list: + lora_info = lora_str.strip().split(",") + lora = Txt2ImgV3LoRA( + model_name=lora_info[0].strip(), + strength=float(lora_info[1]), + ) + result_list.append(lora) + + return result_list + + def _extract_embeddings(self, embeddings_str: str): + if not embeddings_str: + return [] + + embeddings_ori_list = embeddings_str.strip().split(";") + result_list = [] + for embedding_str in embeddings_ori_list: + embedding = Txt2ImgV3Embedding(model_name=embedding_str.strip()) + result_list.append(embedding) + + return result_list + + def _extract_hires_fix(self, hires_fix_str: str): + hires_fix_info = hires_fix_str.strip().split(",") + if "upscaler" in hires_fix_info: + hires_fix = Txt2ImgV3HiresFix( + target_width=int(hires_fix_info[0]), + target_height=int(hires_fix_info[1]), + strength=float(hires_fix_info[2]), + upscaler=hires_fix_info[3].strip(), + ) + else: + hires_fix = Txt2ImgV3HiresFix( + target_width=int(hires_fix_info[0]), + target_height=int(hires_fix_info[1]), + strength=float(hires_fix_info[2]), + ) + + return hires_fix + + def _extract_refiner(self, switch_at: str): + refiner = Txt2ImgV3Refiner(switch_at=float(switch_at)) + return refiner + + def _is_hit_nsfw_detection(self, image: V3TaskImage, confidence_threshold: float) -> bool: + """ + is hit nsfw + """ + if image.nsfw_detection_result is None: + return False + if image.nsfw_detection_result.valid and image.nsfw_detection_result.confidence >= confidence_threshold: + return True + return False diff --git a/api/core/tools/provider/builtin/novitaai/novitaai.py b/api/core/tools/provider/builtin/novitaai/novitaai.py new file mode 100644 index 0000000000000000000000000000000000000000..d5e32eff29373a996eda3797f4b88efdc30f5020 --- /dev/null +++ b/api/core/tools/provider/builtin/novitaai/novitaai.py @@ -0,0 +1,34 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.novitaai.tools.novitaai_txt2img import NovitaAiTxt2ImgTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class NovitaAIProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + result = ( + NovitaAiTxt2ImgTool() + .fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ) + .invoke( + user_id="", + tool_parameters={ + "model_name": "cinenautXLATRUE_cinenautV10_392434.safetensors", + "prompt": "a futuristic city with flying cars", + "negative_prompt": "", + "width": 128, + "height": 128, + "image_num": 1, + "guidance_scale": 7.5, + "seed": -1, + "steps": 1, + }, + ) + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/novitaai/novitaai.yaml b/api/core/tools/provider/builtin/novitaai/novitaai.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3eed8a889c1bd80520f3af0dcd75519c7ca662fb --- /dev/null +++ b/api/core/tools/provider/builtin/novitaai/novitaai.yaml @@ -0,0 +1,32 @@ +identity: + author: Xiao Ley + name: novitaai + label: + en_US: Novita AI + zh_Hans: Novita AI + pt_BR: Novita AI + description: + en_US: Innovative AI for Image Generation + zh_Hans: 用于图像生成的创新人工智能。 + pt_BR: Innovative AI for Image Generation + icon: icon.ico + tags: + - image + - productivity +credentials_for_provider: + api_key: + type: secret-input + required: true + label: + en_US: API Key + zh_Hans: API 密钥 + pt_BR: Chave API + placeholder: + en_US: Please enter your Novita AI API key + zh_Hans: 请输入你的 Novita AI API 密钥 + pt_BR: Por favor, insira sua chave de API do Novita AI + help: + en_US: Get your Novita AI API key from Novita AI + zh_Hans: 从 Novita AI 获取您的 Novita AI API 密钥 + pt_BR: Obtenha sua chave de API do Novita AI na Novita AI + url: https://novita.ai diff --git a/api/core/tools/provider/builtin/novitaai/tools/novitaai_createtile.py b/api/core/tools/provider/builtin/novitaai/tools/novitaai_createtile.py new file mode 100644 index 0000000000000000000000000000000000000000..097b234bd50640668c98ca2d6c1f3c4b66e58209 --- /dev/null +++ b/api/core/tools/provider/builtin/novitaai/tools/novitaai_createtile.py @@ -0,0 +1,54 @@ +from base64 import b64decode +from copy import deepcopy +from typing import Any, Union + +from novita_client import ( # type: ignore + NovitaClient, +) + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.tool.builtin_tool import BuiltinTool + + +class NovitaAiCreateTileTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + if "api_key" not in self.runtime.credentials or not self.runtime.credentials.get("api_key"): + raise ToolProviderCredentialValidationError("Novita AI API Key is required.") + + api_key = self.runtime.credentials.get("api_key") + + client = NovitaClient(api_key=api_key) + param = self._process_parameters(tool_parameters) + client_result = client.create_tile(**param) + + results = [] + results.append( + self.create_blob_message( + blob=b64decode(client_result.image_file), + meta={"mime_type": f"image/{client_result.image_type}"}, + save_as=self.VariableKey.IMAGE.value, + ) + ) + + return results + + def _process_parameters(self, parameters: dict[str, Any]) -> dict[str, Any]: + """ + process parameters + """ + res_parameters = deepcopy(parameters) + + # delete none and empty + keys_to_delete = [k for k, v in res_parameters.items() if v is None or v == ""] + for k in keys_to_delete: + del res_parameters[k] + + return res_parameters diff --git a/api/core/tools/provider/builtin/novitaai/tools/novitaai_createtile.yaml b/api/core/tools/provider/builtin/novitaai/tools/novitaai_createtile.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8e5df5042937d387a17534417de0fa3e70bf42c2 --- /dev/null +++ b/api/core/tools/provider/builtin/novitaai/tools/novitaai_createtile.yaml @@ -0,0 +1,80 @@ +identity: + name: novitaai_createtile + author: Xiao Ley + label: + en_US: Novita AI Create Tile + zh_Hans: Novita AI 创建平铺图案 +description: + human: + en_US: This feature produces images designed for seamless tiling, ideal for creating continuous patterns in fabrics, wallpapers, and various textures. + zh_Hans: 该功能生成设计用于无缝平铺的图像,非常适合用于制作连续图案的织物、壁纸和各种纹理。 + llm: A tool for create images designed for seamless tiling, ideal for creating continuous patterns in fabrics, wallpapers, and various textures. +parameters: + - name: prompt + type: string + required: true + label: + en_US: prompt + zh_Hans: 提示 + human_description: + en_US: Positive prompt word of the created tile, divided by `,`, Range [1, 512]. Only English input is allowed. + zh_Hans: 生成平铺图案的正向提示,用 `,` 分隔,范围 [1, 512]。仅允许输入英文。 + llm_description: Image prompt of Novita AI, you should describe the image you want to generate as a list of words as possible as detailed, divided by `,`, Range [1, 512]. Only English input is allowed. + form: llm + - name: negative_prompt + type: string + required: false + label: + en_US: negative prompt + zh_Hans: 负向提示 + human_description: + en_US: Negtive prompt word of the created tile, divided by `,`, Range [1, 512]. Only English input is allowed. + zh_Hans: 生成平铺图案的负向提示,用 `,` 分隔,范围 [1, 512]。仅允许输入英文。 + llm_description: Image negative prompt of Novita AI, divided by `,`, Range [1, 512]. Only English input is allowed. + form: llm + - name: width + type: number + default: 256 + min: 128 + max: 1024 + required: true + label: + en_US: width + zh_Hans: 宽 + human_description: + en_US: Image width, Range [128, 1024]. + zh_Hans: 图像宽度,范围 [128, 1024] + form: form + - name: height + type: number + default: 256 + min: 128 + max: 1024 + required: true + label: + en_US: height + zh_Hans: 高 + human_description: + en_US: Image height, Range [128, 1024]. + zh_Hans: 图像高度,范围 [128, 1024] + form: form + - name: response_image_type + type: select + default: jpeg + required: false + label: + en_US: response image type + zh_Hans: 响应图像类型 + human_description: + en_US: Response image type, png or jpeg + zh_Hans: 响应图像类型,png 或 jpeg + form: form + options: + - value: jpeg + label: + en_US: jpeg + zh_Hans: jpeg + - value: png + label: + en_US: png + zh_Hans: png diff --git a/api/core/tools/provider/builtin/novitaai/tools/novitaai_modelquery.py b/api/core/tools/provider/builtin/novitaai/tools/novitaai_modelquery.py new file mode 100644 index 0000000000000000000000000000000000000000..a200ee81231f003f26bb43e8c2e8c12910a3552c --- /dev/null +++ b/api/core/tools/provider/builtin/novitaai/tools/novitaai_modelquery.py @@ -0,0 +1,148 @@ +import json +from copy import deepcopy +from typing import Any, Union + +from pandas import DataFrame +from yarl import URL + +from core.helper import ssrf_proxy +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.tool.builtin_tool import BuiltinTool + + +class NovitaAiModelQueryTool(BuiltinTool): + _model_query_endpoint = "https://api.novita.ai/v3/model" + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + if "api_key" not in self.runtime.credentials or not self.runtime.credentials.get("api_key"): + raise ToolProviderCredentialValidationError("Novita AI API Key is required.") + + api_key = self.runtime.credentials.get("api_key") + headers = {"Content-Type": "application/json", "Authorization": "Bearer " + api_key} + params = self._process_parameters(tool_parameters) + result_type = params.get("result_type") + del params["result_type"] + + models_data = self._query_models( + models_data=[], + headers=headers, + params=params, + recursive=result_type not in {"first sd_name", "first name sd_name pair"}, + ) + + result_str = "" + if result_type == "first sd_name": + result_str = models_data[0]["sd_name_in_api"] if len(models_data) > 0 else "" + elif result_type == "first name sd_name pair": + result_str = ( + json.dumps({"name": models_data[0]["name"], "sd_name": models_data[0]["sd_name_in_api"]}) + if len(models_data) > 0 + else "" + ) + elif result_type == "sd_name array": + sd_name_array = [model["sd_name_in_api"] for model in models_data] if len(models_data) > 0 else [] + result_str = json.dumps(sd_name_array) + elif result_type == "name array": + name_array = [model["name"] for model in models_data] if len(models_data) > 0 else [] + result_str = json.dumps(name_array) + elif result_type == "name sd_name pair array": + name_sd_name_pair_array = ( + [{"name": model["name"], "sd_name": model["sd_name_in_api"]} for model in models_data] + if len(models_data) > 0 + else [] + ) + result_str = json.dumps(name_sd_name_pair_array) + elif result_type == "whole info array": + result_str = json.dumps(models_data) + else: + raise NotImplementedError + + return self.create_text_message(result_str) + + def _query_models( + self, + models_data: list, + headers: dict[str, Any], + params: dict[str, Any], + pagination_cursor: str = "", + recursive: bool = True, + ) -> list: + """ + query models + """ + inside_params = deepcopy(params) + + if pagination_cursor != "": + inside_params["pagination.cursor"] = pagination_cursor + + response = ssrf_proxy.get( + url=str(URL(self._model_query_endpoint)), headers=headers, params=params, timeout=(10, 60) + ) + + res_data = response.json() + + models_data.extend(res_data["models"]) + + res_data_len = len(res_data["models"]) + if res_data_len == 0 or res_data_len < int(params["pagination.limit"]) or recursive is False: + # deduplicate + df = DataFrame.from_dict(models_data) + df_unique = df.drop_duplicates(subset=["id"]) + models_data = df_unique.to_dict("records") + return models_data + + return self._query_models( + models_data=models_data, + headers=headers, + params=inside_params, + pagination_cursor=res_data["pagination"]["next_cursor"], + ) + + def _process_parameters(self, parameters: dict[str, Any]) -> dict[str, Any]: + """ + process parameters + """ + process_parameters = deepcopy(parameters) + res_parameters = {} + + # delete none or empty + keys_to_delete = [k for k, v in process_parameters.items() if v is None or v == ""] + for k in keys_to_delete: + del process_parameters[k] + + if "query" in process_parameters and process_parameters.get("query") != "unspecified": + res_parameters["filter.query"] = process_parameters["query"] + + if "visibility" in process_parameters and process_parameters.get("visibility") != "unspecified": + res_parameters["filter.visibility"] = process_parameters["visibility"] + + if "source" in process_parameters and process_parameters.get("source") != "unspecified": + res_parameters["filter.source"] = process_parameters["source"] + + if "type" in process_parameters and process_parameters.get("type") != "unspecified": + res_parameters["filter.types"] = process_parameters["type"] + + if "is_sdxl" in process_parameters: + if process_parameters["is_sdxl"] == "true": + res_parameters["filter.is_sdxl"] = True + elif process_parameters["is_sdxl"] == "false": + res_parameters["filter.is_sdxl"] = False + + res_parameters["result_type"] = process_parameters.get("result_type", "first sd_name") + + res_parameters["pagination.limit"] = ( + 1 + if res_parameters.get("result_type") == "first sd_name" + or res_parameters.get("result_type") == "first name sd_name pair" + else 100 + ) + + return res_parameters diff --git a/api/core/tools/provider/builtin/novitaai/tools/novitaai_modelquery.yaml b/api/core/tools/provider/builtin/novitaai/tools/novitaai_modelquery.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a14795e45e0e4f200839f191f6b3b5b2b47d6dd6 --- /dev/null +++ b/api/core/tools/provider/builtin/novitaai/tools/novitaai_modelquery.yaml @@ -0,0 +1,175 @@ +identity: + name: novitaai_modelquery + author: Xiao Ley + label: + en_US: Novita AI Model Query + zh_Hans: Novita AI 模型查询 +description: + human: + en_US: Retrieve information on both public and private models. It allows users to access details such as model specifications, status, and usage guidelines, ensuring comprehensive insight into the available modeling resources. + zh_Hans: 检索公开和私有模型信息。它允许用户访问模型规范、状态和使用指南等详细信息,确保了解可用的建模资源。 + llm: A tool for retrieve information on both public and private Novita AI models. +parameters: + - name: query + type: string + required: false + label: + en_US: query + zh_Hans: 查询 + human_description: + en_US: Seaching the content of sd_name, name, tags. + zh_Hans: 搜索 sd_name、name、tags 中的内容 + llm_description: Enter the content to search + form: llm + - name: result_type + type: select + default: "first sd_name" + required: true + label: + en_US: result format + zh_Hans: 结果格式 + human_description: + en_US: The format of result + zh_Hans: 请求结果的格式 + form: form + options: + - value: "first sd_name" + label: + en_US: "first sd_name" + zh_Hans: "第一个 sd_name" + - value: "first name sd_name pair" + label: + en_US: "first name and sd_name pair: {name, sd_name}" + zh_Hans: "第一个 name sd_name 组合:{name, sd_name}" + - value: "sd_name array" + label: + en_US: "sd_name array: [sd_name]" + zh_Hans: "sd_name 数组:[sd_name]" + - value: "name array" + label: + en_US: "name array: [name]" + zh_Hans: "name 数组:[name]" + - value: "name sd_name pair array" + label: + en_US: "name and sd_name pair array: [{name, sd_name}]" + zh_Hans: "name sd_name 组合数组:[{name, sd_name}]" + - value: "whole info array" + label: + en_US: whole info array + zh_Hans: 完整信息数组 + - name: visibility + type: select + default: unspecified + required: false + label: + en_US: visibility + zh_Hans: 可见性 + human_description: + en_US: Whether the model is public or private + zh_Hans: 模型是否公开或私有 + form: form + options: + - value: unspecified + label: + en_US: Unspecified + zh_Hans: 未指定 + - value: public + label: + en_US: Public + zh_Hans: 公开 + - value: private + label: + en_US: Private + zh_Hans: 私有 + - name: source + type: select + default: unspecified + required: false + label: + en_US: source + zh_Hans: 来源 + human_description: + en_US: Source of the model + zh_Hans: 模型来源 + form: form + options: + - value: unspecified + label: + en_US: Unspecified + zh_Hans: 未指定 + - value: civitai + label: + en_US: Civitai + zh_Hans: Civitai + - value: training + label: + en_US: Training + zh_Hans: 训练 + - value: uploading + label: + en_US: Uploading + zh_Hans: 上传 + - name: type + type: select + default: unspecified + required: false + label: + en_US: type + zh_Hans: 类型 + human_description: + en_US: Specifies the type of models to include in the query. + zh_Hans: 指定要查询的模型类型 + form: form + options: + - value: unspecified + label: + en_US: Unspecified + zh_Hans: 未指定 + - value: checkpoint + label: + en_US: Checkpoint + zh_Hans: Checkpoint + - value: lora + label: + en_US: LoRA + zh_Hans: LoRA + - value: vae + label: + en_US: VAE + zh_Hans: VAE + - value: controlnet + label: + en_US: ControlNet + zh_Hans: ControlNet + - value: upscaler + label: + en_US: Upscaler + zh_Hans: Upscaler + - value: textualinversion + label: + en_US: Textual inversion + zh_Hans: Textual Inversion + - name: is_sdxl + type: select + default: unspecified + required: false + label: + en_US: is sdxl + zh_Hans: 是否是 SDXL + human_description: + en_US: Whether sdxl model or not. Setting this parameter to `true` includes only sdxl models in the query results, which are typically large-scale, high-performance models designed for extensive data processing tasks. Conversely, setting it to `false` excludes these models from the results. If left unspecified, the filter will not discriminate based on the sdxl classification, including all model types in the search results. + zh_Hans: 是否是 SDXL 模型。设置此参数为 `是`,只查询 SDXL 模型,并包含大规模,高性能的模型。相反,设置为 `否`,将排除这些模型。如果未指定,将不会根据 SDXL 分类进行区分,包括查询结果中的所有模型类型。 + form: form + options: + - value: unspecified + label: + en_US: Unspecified + zh_Hans: 未指定 + - value: "true" + label: + en_US: "True" + zh_Hans: 是 + - value: "false" + label: + en_US: "False" + zh_Hans: 否 diff --git a/api/core/tools/provider/builtin/novitaai/tools/novitaai_txt2img.py b/api/core/tools/provider/builtin/novitaai/tools/novitaai_txt2img.py new file mode 100644 index 0000000000000000000000000000000000000000..297a27abba667a3c5f15cc15942282998e1bda33 --- /dev/null +++ b/api/core/tools/provider/builtin/novitaai/tools/novitaai_txt2img.py @@ -0,0 +1,90 @@ +from base64 import b64decode +from copy import deepcopy +from typing import Any, Union + +from novita_client import ( # type: ignore + NovitaClient, +) + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.novitaai._novita_tool_base import NovitaAiToolBase +from core.tools.tool.builtin_tool import BuiltinTool + + +class NovitaAiTxt2ImgTool(BuiltinTool, NovitaAiToolBase): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + if "api_key" not in self.runtime.credentials or not self.runtime.credentials.get("api_key"): + raise ToolProviderCredentialValidationError("Novita AI API Key is required.") + + api_key = self.runtime.credentials.get("api_key") + + client = NovitaClient(api_key=api_key) + param = self._process_parameters(tool_parameters) + client_result = client.txt2img_v3(**param) + + results = [] + for image_encoded, image in zip(client_result.images_encoded, client_result.images): + if self._is_hit_nsfw_detection(image, 0.8): + results = self.create_text_message(text="NSFW detected!") + break + + results.append( + self.create_blob_message( + blob=b64decode(image_encoded), + meta={"mime_type": f"image/{image.image_type}"}, + save_as=self.VariableKey.IMAGE.value, + ) + ) + + return results + + def _process_parameters(self, parameters: dict[str, Any]) -> dict[str, Any]: + """ + process parameters + """ + res_parameters = deepcopy(parameters) + + # delete none and empty + keys_to_delete = [k for k, v in res_parameters.items() if v is None or v == ""] + for k in keys_to_delete: + del res_parameters[k] + + if "clip_skip" in res_parameters and res_parameters.get("clip_skip") == 0: + del res_parameters["clip_skip"] + + if "refiner_switch_at" in res_parameters and res_parameters.get("refiner_switch_at") == 0: + del res_parameters["refiner_switch_at"] + + if "enabled_enterprise_plan" in res_parameters: + res_parameters["enterprise_plan"] = {"enabled": res_parameters["enabled_enterprise_plan"]} + del res_parameters["enabled_enterprise_plan"] + + if "nsfw_detection_level" in res_parameters: + res_parameters["nsfw_detection_level"] = int(res_parameters["nsfw_detection_level"]) + + # process loras + if "loras" in res_parameters: + res_parameters["loras"] = self._extract_loras(res_parameters.get("loras")) + + # process embeddings + if "embeddings" in res_parameters: + res_parameters["embeddings"] = self._extract_embeddings(res_parameters.get("embeddings")) + + # process hires_fix + if "hires_fix" in res_parameters: + res_parameters["hires_fix"] = self._extract_hires_fix(res_parameters.get("hires_fix")) + + # process refiner + if "refiner_switch_at" in res_parameters: + res_parameters["refiner"] = self._extract_refiner(res_parameters.get("refiner_switch_at")) + del res_parameters["refiner_switch_at"] + + return res_parameters diff --git a/api/core/tools/provider/builtin/novitaai/tools/novitaai_txt2img.yaml b/api/core/tools/provider/builtin/novitaai/tools/novitaai_txt2img.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d625a643f915b1402fb0cb29d831328085c360b7 --- /dev/null +++ b/api/core/tools/provider/builtin/novitaai/tools/novitaai_txt2img.yaml @@ -0,0 +1,341 @@ +identity: + name: novitaai_txt2img + author: Xiao Ley + label: + en_US: Novita AI Text to Image + zh_Hans: Novita AI 文字转图像 +description: + human: + en_US: Generate images from text prompts using Stable Diffusion models + zh_Hans: 通过 Stable Diffusion 模型根据文字提示生成图像 + llm: A tool for generate images from English text prompts. +parameters: + - name: model_name + type: string + required: true + label: + en_US: model name + zh_Hans: 模块名字 + human_description: + en_US: Specify the name of the model checkpoint. You can use the "Novita AI Model Query" tool to query the corresponding "sd_name" value (type select "Checkpoint"). + zh_Hans: 指定 Model Checkpoint 名称。可通过“Novita AI 模型请求”工具查询对应的“sd_name”值(类型选择“Checkpoint”)。 + form: form + - name: prompt + type: string + required: true + label: + en_US: prompt + zh_Hans: 提示 + human_description: + en_US: Text input required to guide the image generation, divided by `,`, Range [1, 1024]. Only English input is allowed. + zh_Hans: 生成图像的正向提示,用 `,` 分隔,范围 [1, 1024]。仅允许输入英文。 + llm_description: Image prompt of Novita AI, you should describe the image you want to generate as a list of words as possible as detailed, divided by `,`, Range [1, 1024]. Only English input is allowed. + form: llm + - name: negative_prompt + type: string + required: false + label: + en_US: negative prompt + zh_Hans: 负向提示 + human_description: + en_US: Text input that will not guide the image generation, divided by `,`, Range [1, 1024]. Only English input is allowed. + zh_Hans: 生成图像的负向提示,用 `,` 分隔,范围 [1, 1024]。仅允许输入英文。 + llm_description: Image negative prompt of Novita AI, divided by `,`, Range [1, 1024]. Only English input is allowed. + form: llm + - name: width + type: number + default: 512 + min: 128 + max: 2048 + required: true + label: + en_US: width + zh_Hans: 宽 + human_description: + en_US: Image width, Range [128, 2048]. + zh_Hans: 图像宽度,范围 [128, 2048] + form: form + - name: height + type: number + default: 512 + min: 128 + max: 2048 + required: true + label: + en_US: height + zh_Hans: 高 + human_description: + en_US: Image height, Range [128, 2048]. + zh_Hans: 图像高度,范围 [128, 2048] + form: form + - name: image_num + type: number + default: 1 + min: 1 + max: 8 + required: true + label: + en_US: image num + zh_Hans: 图片数 + human_description: + en_US: Image num, Range [1, 8]. + zh_Hans: 图片数,范围 [1, 8] + form: form + - name: steps + type: number + default: 20 + min: 1 + max: 100 + required: true + label: + en_US: steps + zh_Hans: 步数 + human_description: + en_US: The number of denoising steps. More steps usually can produce higher quality images, but take more time to generate, Range [1, 100]. + zh_Hans: 生成步数。更多步数可能会产生更好的图像,但生成时间更长,范围 [1, 100] + form: form + - name: seed + type: number + default: -1 + required: true + label: + en_US: seed + zh_Hans: 种子 + human_description: + en_US: A seed is a number from which Stable Diffusion generates noise, which, makes generation deterministic. Using the same seed and set of parameters will produce identical image each time, minimum -1. + zh_Hans: 种子是 Stable Diffusion 生成噪声的数字,它使生成具有确定性。使用相同的种子和参数设置将生成每次生成相同的图像,最小值 -1。 + form: form + - name: clip_skip + type: number + min: 1 + max: 12 + required: false + label: + en_US: clip skip + zh_Hans: 层跳过数 + human_description: + en_US: This parameter indicates the number of layers to stop from the bottom during optimization, so clip_skip on 2 would mean, that in SD1.x model where the CLIP has 12 layers, you would stop at 10th layer, Range [1, 12], get reference at https://novita.ai/get-started/Misc.html#what-s-clip-skip. + zh_Hans: 此参数表示优化过程中从底部停止的层数,因此 clip_skip 的值为 2,表示在 SD1.x 模型中,CLIP 有 12 层,你将停止在 10 层,范围 [1, 12],参考 https://novita.ai/get-started/Misc.html#what-s-clip-skip。 + form: form + - name: guidance_scale + type: number + default: "7.5" + min: 1.0 + max: 30.0 + required: true + label: + en_US: guidance scale + zh_Hans: 提示词遵守程度 + human_description: + en_US: This setting says how close the Stable Diffusion will listen to your prompt, higer guidance forces the model to better follow the prompt, but result in lower quality output.Range [1, 30]. + zh_Hans: 此设置表明 Stable Diffusion 如何听从您的提示,较高的 guidance_scale 会强制模型更好跟随提示,但结果会更低质量输出。范围 [1.0, 30.0]。 + form: form + - name: sampler_name + type: select + required: true + label: + en_US: sampler name + zh_Hans: 采样器名称 + human_description: + en_US: This parameter determines the denoising algorithm employed during the sampling phase of Stable Diffusion. Get reference at https://novita.ai/get-started/Misc.htmll#what-is-samplers. + zh_Hans: 此参数决定了在稳定扩散采样阶段使用的去噪算法。参考 https://novita.ai/get-started/Misc.htmll#what-is-samplers。 + form: form + options: + - value: "Euler a" + label: + en_US: Euler a + zh_Hans: Euler a + - value: "Euler" + label: + en_US: Euler + zh_Hans: Euler + - value: "LMS" + label: + en_US: LMS + zh_Hans: LMS + - value: "Heun" + label: + en_US: Heun + zh_Hans: Heun + - value: "DPM2" + label: + en_US: DPM2 + zh_Hans: DPM2 + - value: "DPM2 a" + label: + en_US: DPM2 a + zh_Hans: DPM2 a + - value: "DPM++ 2S a" + label: + en_US: DPM++ 2S a + zh_Hans: DPM++ 2S a + - value: "DPM++ 2M" + label: + en_US: DPM++ 2M + zh_Hans: DPM++ 2M + - value: "DPM++ SDE" + label: + en_US: DPM++ SDE + zh_Hans: DPM++ SDE + - value: "DPM fast" + label: + en_US: DPM fast + zh_Hans: DPM fast + - value: "DPM adaptive" + label: + en_US: DPM adaptive + zh_Hans: DPM adaptive + - value: "LMS Karras" + label: + en_US: LMS Karras + zh_Hans: LMS Karras + - value: "DPM2 Karras" + label: + en_US: DPM2 Karras + zh_Hans: DPM2 Karras + - value: "DPM2 a Karras" + label: + en_US: DPM2 a Karras + zh_Hans: DPM2 a Karras + - value: "DPM++ 2S a Karras" + label: + en_US: DPM++ 2S a Karras + zh_Hans: DPM++ 2S a Karras + - value: "DPM++ 2M Karras" + label: + en_US: DPM++ 2M Karras + zh_Hans: DPM++ 2M Karras + - value: "DPM++ SDE Karras" + label: + en_US: DPM++ SDE Karras + zh_Hans: DPM++ SDE Karras + - value: "DDIM" + label: + en_US: DDIM + zh_Hans: DDIM + - value: "PLMS" + label: + en_US: PLMS + zh_Hans: PLMS + - value: "UniPC" + label: + en_US: UniPC + zh_Hans: UniPC + - name: sd_vae + type: string + required: false + label: + en_US: sd vae + zh_Hans: sd vae + human_description: + en_US: VAE(Variational Autoencoder), get reference at https://novita.ai/get-started/Misc.html#what-s-variational-autoencoders-vae. You can use the "Novita AI Model Query" tool to query the corresponding "sd_name" value (type select "VAE"). + zh_Hans: VAE(变分自编码器),参考 https://novita.ai/get-started/Misc.html#what-s-variational-autoencoders-vae。可通过“Novita AI 模型请求”工具查询对应的“sd_name”值(类型选择“VAE”)。 + form: form + - name: loras + type: string + required: false + label: + en_US: loRAs + zh_Hans: loRAs + human_description: + en_US: LoRA models. Currenlty supports up to 5 LoRAs. You can use the "Novita AI Model Query" tool to query the corresponding "sd_name" value (type select "LoRA"). Input template is ",;,;...". Such as"Film Grain style_331903,0.5;DoggystylePOV_9600,0.5" + zh_Hans: LoRA 模型。目前仅支持 5 个 LoRA。可通过“Novita AI 模型请求”工具查询对应的“sd_name”值(类型选择“LoRA”)。输入模板:“,;,;...”,例如:“Film Grain style_331903,0.5;DoggystylePOV_9600,0.5” + form: form + - name: embeddings + type: string + required: false + label: + en_US: text embeddings + zh_Hans: 文本嵌入 + human_description: + en_US: Textual Inversion is a training method for personalizing models by learning new text embeddings from a few example images, currenlty supports up to 5 embeddings. You can use the "Novita AI Model Query" tool to query the corresponding "sd_name" value (type select "Text Inversion"). Input template is ";;...". Such as "EasyNegativeV2_75525;AS-YoungerV2" + zh_Hans: 文本反转是一种通过从一些示例图像中学习新的文本嵌入来个性化模型的训练方法,目前仅支持 5 个嵌入。可通过“Novita AI 模型请求”工具查询对应的“sd_name”值(类型选择“Text Inversion”)。输入模板:“;;...”,例如:“EasyNegativeV2_75525;AS-YoungerV2” + form: form + - name: hires_fix + type: string + required: false + label: + en_US: hires fix + zh_Hans: 高分辨率修复 + human_description: + en_US: Use high resolution image fix. Input template is ",,,". Such as "1024,1024,0.8", "1024,1024,0.8,RealESRGAN_x4plus_anime_6B" + zh_Hans: 使用高分辨率修复。输入模板 “,,,”。例如 “1024,1024,0.8”、“1024,1024,0.8,RealESRGAN_x4plus_anime_6B” + form: form + - name: refiner_switch_at + type: number + min: 0.0 + max: 1.0 + required: false + label: + en_US: refiner switch at + zh_Hans: 重采样参与时刻 + human_description: + en_US: This parameter in the context of a refiner allows you to set the extent to which the refiner alters the output of a model. When set to 0, the refiner has no effect; at 1, it's fully active. Intermediate values like 0.5 provide a balanced effect, where the refiner is moderately engaged, enhancing or adjusting the output without dominating the original model's characteristics. This setting is particularly useful for fine-tuning the output to achieve a desired balance between refinement and the original generative features, Range [0, 1.0]. Is not all models support refiners! + zh_Hans: 此参数允许您设置重采样更改模型输出的程度。当设置为0时,重采样不起作用;1时,它处于完全活动状态。像0.5这样的中间值提供了一种平衡效果,其中重采样适度参与,增强或调整输出,而不会主导原始模型的特性。此设置对于微调输出特别有用,范围 [0, 1.0]。不是所有模型都支持重采样! + form: form + - name: response_image_type + type: select + default: jpeg + required: false + label: + en_US: response image type + zh_Hans: 响应图像类型 + human_description: + en_US: Response image type, png or jpeg + zh_Hans: 响应图像类型,png 或 jpeg + form: form + options: + - value: jpeg + label: + en_US: jpeg + zh_Hans: jpeg + - value: png + label: + en_US: png + zh_Hans: png + - name: enabled_enterprise_plan + type: boolean + default: false + required: false + label: + en_US: enterprise plan enabled + zh_Hans: 企业版计划启用 + human_description: + en_US: Enable enterprise plan + zh_Hans: 启用企业版计划 + form: form + - name: enable_nsfw_detection + type: boolean + default: false + required: false + label: + en_US: enable nsfw detection + zh_Hans: 启用 NSFW 检测 + human_description: + en_US: Enable nsfw detection + zh_Hans: 启用 NSFW 检测 + form: form + - name: nsfw_detection_level + type: select + default: "2" + required: false + label: + en_US: nsfw detection level + zh_Hans: NSFW 检测级别 + human_description: + en_US: Nsfw detection level, from low to high + zh_Hans: NSFW 检测级别,越高越严格 + form: form + options: + - value: "0" + label: + en_US: low + zh_Hans: 低 + - value: "1" + label: + en_US: middle + zh_Hans: 中 + - value: "2" + label: + en_US: high + zh_Hans: 高 diff --git a/api/core/tools/provider/builtin/onebot/_assets/icon.ico b/api/core/tools/provider/builtin/onebot/_assets/icon.ico new file mode 100644 index 0000000000000000000000000000000000000000..1b07e965b9910b4b006bc112378a8ba0306895a8 Binary files /dev/null and b/api/core/tools/provider/builtin/onebot/_assets/icon.ico differ diff --git a/api/core/tools/provider/builtin/onebot/onebot.py b/api/core/tools/provider/builtin/onebot/onebot.py new file mode 100644 index 0000000000000000000000000000000000000000..b8e5ed24d6b43f327c98a1a307f6224e32bf5e7a --- /dev/null +++ b/api/core/tools/provider/builtin/onebot/onebot.py @@ -0,0 +1,10 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class OneBotProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + if not credentials.get("ob11_http_url"): + raise ToolProviderCredentialValidationError("OneBot HTTP URL is required.") diff --git a/api/core/tools/provider/builtin/onebot/onebot.yaml b/api/core/tools/provider/builtin/onebot/onebot.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1922adc4de4d56a4110ff3fd59a4ab95da400839 --- /dev/null +++ b/api/core/tools/provider/builtin/onebot/onebot.yaml @@ -0,0 +1,35 @@ +identity: + author: RockChinQ + name: onebot + label: + en_US: OneBot v11 Protocol + zh_Hans: OneBot v11 协议 + description: + en_US: Unofficial OneBot v11 Protocol Tool + zh_Hans: 非官方 OneBot v11 协议工具 + icon: icon.ico +credentials_for_provider: + ob11_http_url: + type: text-input + required: true + label: + en_US: HTTP URL + zh_Hans: HTTP URL + description: + en_US: Forward HTTP URL of OneBot v11 + zh_Hans: OneBot v11 正向 HTTP URL + help: + en_US: Fill this with the HTTP URL of your OneBot server + zh_Hans: 请在你的 OneBot 协议端开启 正向 HTTP 并填写其 URL + access_token: + type: secret-input + required: false + label: + en_US: Access Token + zh_Hans: 访问令牌 + description: + en_US: Access Token for OneBot v11 Protocol + zh_Hans: OneBot 协议访问令牌 + help: + en_US: Fill this if you set a access token in your OneBot server + zh_Hans: 如果你在 OneBot 服务器中设置了 access token,请填写此项 diff --git a/api/core/tools/provider/builtin/onebot/tools/__init__.py b/api/core/tools/provider/builtin/onebot/tools/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/tools/provider/builtin/onebot/tools/send_group_msg.py b/api/core/tools/provider/builtin/onebot/tools/send_group_msg.py new file mode 100644 index 0000000000000000000000000000000000000000..9c95bbc2ae8d2deae9c0fc6aa4c4ec5a545e924e --- /dev/null +++ b/api/core/tools/provider/builtin/onebot/tools/send_group_msg.py @@ -0,0 +1,39 @@ +from typing import Any, Union + +import requests +from yarl import URL + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class SendGroupMsg(BuiltinTool): + """OneBot v11 Tool: Send Group Message""" + + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + # Get parameters + send_group_id = tool_parameters.get("group_id", "") + + message = tool_parameters.get("message", "") + if not message: + return self.create_json_message({"error": "Message is empty."}) + + auto_escape = tool_parameters.get("auto_escape", False) + + try: + url = URL(self.runtime.credentials["ob11_http_url"]) / "send_group_msg" + + resp = requests.post( + url, + json={"group_id": send_group_id, "message": message, "auto_escape": auto_escape}, + headers={"Authorization": "Bearer " + self.runtime.credentials["access_token"]}, + ) + + if resp.status_code != 200: + return self.create_json_message({"error": f"Failed to send group message: {resp.text}"}) + + return self.create_json_message({"response": resp.json()}) + except Exception as e: + return self.create_json_message({"error": f"Failed to send group message: {e}"}) diff --git a/api/core/tools/provider/builtin/onebot/tools/send_group_msg.yaml b/api/core/tools/provider/builtin/onebot/tools/send_group_msg.yaml new file mode 100644 index 0000000000000000000000000000000000000000..64beaa85457a3aaece110974078bd9cb93f5b233 --- /dev/null +++ b/api/core/tools/provider/builtin/onebot/tools/send_group_msg.yaml @@ -0,0 +1,46 @@ +identity: + name: send_group_msg + author: RockChinQ + label: + en_US: Send Group Message + zh_Hans: 发送群消息 +description: + human: + en_US: Send a message to a group + zh_Hans: 发送消息到群聊 + llm: A tool for sending a message segment to a group +parameters: + - name: group_id + type: number + required: true + label: + en_US: Target Group ID + zh_Hans: 目标群 ID + human_description: + en_US: The group ID of the target group + zh_Hans: 目标群的群 ID + llm_description: The group ID of the target group + form: llm + - name: message + type: string + required: true + label: + en_US: Message + zh_Hans: 消息 + human_description: + en_US: The message to send + zh_Hans: 要发送的消息。支持 CQ码(需要同时设置 auto_escape 为 true) + llm_description: The message to send + form: llm + - name: auto_escape + type: boolean + required: false + default: false + label: + en_US: Auto Escape + zh_Hans: 自动转义 + human_description: + en_US: If true, the message will be treated as a CQ code for parsing, otherwise it will be treated as plain text for direct sending. Since Dify currently does not support passing Object-format message chains, developers can send complex message components through CQ codes. + zh_Hans: 若为 true 则会把 message 视为 CQ 码解析,否则视为 纯文本 直接发送。由于 Dify 目前不支持传入 Object格式 的消息,故开发者可以通过 CQ 码来发送复杂消息组件。 + llm_description: If true, the message will be treated as a CQ code for parsing, otherwise it will be treated as plain text for direct sending. + form: form diff --git a/api/core/tools/provider/builtin/onebot/tools/send_private_msg.py b/api/core/tools/provider/builtin/onebot/tools/send_private_msg.py new file mode 100644 index 0000000000000000000000000000000000000000..1174c7f07d002f7cb101d8e0a10840f4c5db9f75 --- /dev/null +++ b/api/core/tools/provider/builtin/onebot/tools/send_private_msg.py @@ -0,0 +1,39 @@ +from typing import Any, Union + +import requests +from yarl import URL + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class SendPrivateMsg(BuiltinTool): + """OneBot v11 Tool: Send Private Message""" + + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + # Get parameters + send_user_id = tool_parameters.get("user_id", "") + + message = tool_parameters.get("message", "") + if not message: + return self.create_json_message({"error": "Message is empty."}) + + auto_escape = tool_parameters.get("auto_escape", False) + + try: + url = URL(self.runtime.credentials["ob11_http_url"]) / "send_private_msg" + + resp = requests.post( + url, + json={"user_id": send_user_id, "message": message, "auto_escape": auto_escape}, + headers={"Authorization": "Bearer " + self.runtime.credentials["access_token"]}, + ) + + if resp.status_code != 200: + return self.create_json_message({"error": f"Failed to send private message: {resp.text}"}) + + return self.create_json_message({"response": resp.json()}) + except Exception as e: + return self.create_json_message({"error": f"Failed to send private message: {e}"}) diff --git a/api/core/tools/provider/builtin/onebot/tools/send_private_msg.yaml b/api/core/tools/provider/builtin/onebot/tools/send_private_msg.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8200ce4a83f4e28dba9a69524a43dc42d146118c --- /dev/null +++ b/api/core/tools/provider/builtin/onebot/tools/send_private_msg.yaml @@ -0,0 +1,46 @@ +identity: + name: send_private_msg + author: RockChinQ + label: + en_US: Send Private Message + zh_Hans: 发送私聊消息 +description: + human: + en_US: Send a private message to a user + zh_Hans: 发送私聊消息给用户 + llm: A tool for sending a message segment to a user in private chat +parameters: + - name: user_id + type: number + required: true + label: + en_US: Target User ID + zh_Hans: 目标用户 ID + human_description: + en_US: The user ID of the target user + zh_Hans: 目标用户的用户 ID + llm_description: The user ID of the target user + form: llm + - name: message + type: string + required: true + label: + en_US: Message + zh_Hans: 消息 + human_description: + en_US: The message to send + zh_Hans: 要发送的消息。支持 CQ码(需要同时设置 auto_escape 为 true) + llm_description: The message to send + form: llm + - name: auto_escape + type: boolean + required: false + default: false + label: + en_US: Auto Escape + zh_Hans: 自动转义 + human_description: + en_US: If true, the message will be treated as a CQ code for parsing, otherwise it will be treated as plain text for direct sending. Since Dify currently does not support passing Object-format message chains, developers can send complex message components through CQ codes. + zh_Hans: 若为 true 则会把 message 视为 CQ 码解析,否则视为 纯文本 直接发送。由于 Dify 目前不支持传入 Object格式 的消息,故开发者可以通过 CQ 码来发送复杂消息组件。 + llm_description: If true, the message will be treated as a CQ code for parsing, otherwise it will be treated as plain text for direct sending. + form: form diff --git a/api/core/tools/provider/builtin/openweather/_assets/icon.svg b/api/core/tools/provider/builtin/openweather/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..f06cd87e64c9d3bf2104f02d919307a71e947cdc --- /dev/null +++ b/api/core/tools/provider/builtin/openweather/_assets/icon.svg @@ -0,0 +1,12 @@ + + + + + + + + + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/openweather/openweather.py b/api/core/tools/provider/builtin/openweather/openweather.py new file mode 100644 index 0000000000000000000000000000000000000000..9e40249aba6b40e4cd778ea0febdbaa389a8c2a4 --- /dev/null +++ b/api/core/tools/provider/builtin/openweather/openweather.py @@ -0,0 +1,29 @@ +import requests + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +def query_weather(city="Beijing", units="metric", language="zh_cn", api_key=None): + url = "https://api.openweathermap.org/data/2.5/weather" + params = {"q": city, "appid": api_key, "units": units, "lang": language} + + return requests.get(url, params=params) + + +class OpenweatherProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + try: + if "api_key" not in credentials or not credentials.get("api_key"): + raise ToolProviderCredentialValidationError("Open weather API key is required.") + apikey = credentials.get("api_key") + try: + response = query_weather(api_key=apikey) + if response.status_code == 200: + pass + else: + raise ToolProviderCredentialValidationError((response.json()).get("info")) + except Exception as e: + raise ToolProviderCredentialValidationError("Open weather API Key is invalid. {}".format(e)) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/openweather/openweather.yaml b/api/core/tools/provider/builtin/openweather/openweather.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d4b66f87f908c6337efd4e93e00e83d980474321 --- /dev/null +++ b/api/core/tools/provider/builtin/openweather/openweather.yaml @@ -0,0 +1,31 @@ +identity: + author: Onelevenvy + name: openweather + label: + en_US: Open weather query + zh_Hans: Open Weather + pt_BR: Consulta de clima open weather + description: + en_US: Weather query toolkit based on Open Weather + zh_Hans: 基于open weather的天气查询工具包 + pt_BR: Kit de consulta de clima baseado no Open Weather + icon: icon.svg + tags: + - weather +credentials_for_provider: + api_key: + type: secret-input + required: true + label: + en_US: API Key + zh_Hans: API Key + pt_BR: Fogo a chave + placeholder: + en_US: Please enter your open weather API Key + zh_Hans: 请输入你的open weather API Key + pt_BR: Insira sua chave de API open weather + help: + en_US: Get your API Key from open weather + zh_Hans: 从open weather获取您的 API Key + pt_BR: Obtenha sua chave de API do open weather + url: https://openweathermap.org diff --git a/api/core/tools/provider/builtin/openweather/tools/weather.py b/api/core/tools/provider/builtin/openweather/tools/weather.py new file mode 100644 index 0000000000000000000000000000000000000000..ed4ec487fa984a11db802add1620599e6616331e --- /dev/null +++ b/api/core/tools/provider/builtin/openweather/tools/weather.py @@ -0,0 +1,52 @@ +import json +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class OpenweatherTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + city = tool_parameters.get("city", "") + if not city: + return self.create_text_message("Please tell me your city") + if "api_key" not in self.runtime.credentials or not self.runtime.credentials.get("api_key"): + return self.create_text_message("OpenWeather API key is required.") + + units = tool_parameters.get("units", "metric") + lang = tool_parameters.get("lang", "zh_cn") + try: + # request URL + url = "https://api.openweathermap.org/data/2.5/weather" + + # request params + params = { + "q": city, + "appid": self.runtime.credentials.get("api_key"), + "units": units, + "lang": lang, + } + response = requests.get(url, params=params) + + if response.status_code == 200: + data = response.json() + return self.create_text_message( + self.summary(user_id=user_id, content=json.dumps(data, ensure_ascii=False)) + ) + else: + error_message = { + "error": f"failed:{response.status_code}", + "data": response.text, + } + # return error + return json.dumps(error_message) + + except Exception as e: + return self.create_text_message("Openweather API Key is invalid. {}".format(e)) diff --git a/api/core/tools/provider/builtin/openweather/tools/weather.yaml b/api/core/tools/provider/builtin/openweather/tools/weather.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f2dae5c2df9c08e21d142715995178184c572b4f --- /dev/null +++ b/api/core/tools/provider/builtin/openweather/tools/weather.yaml @@ -0,0 +1,80 @@ +identity: + name: weather + author: Onelevenvy + label: + en_US: Open Weather Query + zh_Hans: 天气查询 + pt_BR: Previsão do tempo + icon: icon.svg +description: + human: + en_US: Weather forecast inquiry + zh_Hans: 天气查询 + pt_BR: Inquérito sobre previsão meteorológica + llm: A tool when you want to ask about the weather or weather-related question +parameters: + - name: city + type: string + required: true + label: + en_US: city + zh_Hans: 城市 + pt_BR: cidade + human_description: + en_US: Target city for weather forecast query + zh_Hans: 天气预报查询的目标城市 + pt_BR: Cidade de destino para consulta de previsão do tempo + llm_description: If you don't know you can extract the city name from the + question or you can reply:Please tell me your city. You have to extract + the Chinese city name from the question.If the input region is in Chinese + characters for China, it should be replaced with the corresponding English + name, such as '北京' for correct input is 'Beijing' + form: llm + - name: lang + type: select + required: true + human_description: + en_US: language + zh_Hans: 语言 + pt_BR: language + label: + en_US: language + zh_Hans: 语言 + pt_BR: language + form: form + options: + - value: zh_cn + label: + en_US: cn + zh_Hans: 中国 + pt_BR: cn + - value: en_us + label: + en_US: usa + zh_Hans: 美国 + pt_BR: usa + default: zh_cn + - name: units + type: select + required: true + human_description: + en_US: units for temperature + zh_Hans: 温度单位 + pt_BR: units for temperature + label: + en_US: units + zh_Hans: 单位 + pt_BR: units + form: form + options: + - value: metric + label: + en_US: metric + zh_Hans: ℃ + pt_BR: metric + - value: imperial + label: + en_US: imperial + zh_Hans: ℉ + pt_BR: imperial + default: metric diff --git a/api/core/tools/provider/builtin/perplexity/_assets/icon.svg b/api/core/tools/provider/builtin/perplexity/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..c2974c142fc6226be3cf76d6d2042d7b384ebb6c --- /dev/null +++ b/api/core/tools/provider/builtin/perplexity/_assets/icon.svg @@ -0,0 +1,3 @@ + + + diff --git a/api/core/tools/provider/builtin/perplexity/perplexity.py b/api/core/tools/provider/builtin/perplexity/perplexity.py new file mode 100644 index 0000000000000000000000000000000000000000..80518853fb4a4be010e729000f884dbc467a735f --- /dev/null +++ b/api/core/tools/provider/builtin/perplexity/perplexity.py @@ -0,0 +1,38 @@ +from typing import Any + +import requests + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.perplexity.tools.perplexity_search import PERPLEXITY_API_URL +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class PerplexityProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + headers = { + "Authorization": f"Bearer {credentials.get('perplexity_api_key')}", + "Content-Type": "application/json", + } + + payload = { + "model": "llama-3.1-sonar-small-128k-online", + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello"}, + ], + "max_tokens": 5, + "temperature": 0.1, + "top_p": 0.9, + "stream": False, + } + + try: + response = requests.post(PERPLEXITY_API_URL, json=payload, headers=headers) + response.raise_for_status() + except requests.RequestException as e: + raise ToolProviderCredentialValidationError(f"Failed to validate Perplexity API key: {str(e)}") + + if response.status_code != 200: + raise ToolProviderCredentialValidationError( + f"Perplexity API key is invalid. Status code: {response.status_code}" + ) diff --git a/api/core/tools/provider/builtin/perplexity/perplexity.yaml b/api/core/tools/provider/builtin/perplexity/perplexity.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c0b504f300c45a74b7724bec0c4dec275bb53ce0 --- /dev/null +++ b/api/core/tools/provider/builtin/perplexity/perplexity.yaml @@ -0,0 +1,26 @@ +identity: + author: Dify + name: perplexity + label: + en_US: Perplexity + zh_Hans: Perplexity + description: + en_US: Perplexity.AI + zh_Hans: Perplexity.AI + icon: icon.svg + tags: + - search +credentials_for_provider: + perplexity_api_key: + type: secret-input + required: true + label: + en_US: Perplexity API key + zh_Hans: Perplexity API key + placeholder: + en_US: Please input your Perplexity API key + zh_Hans: 请输入你的 Perplexity API key + help: + en_US: Get your Perplexity API key from Perplexity + zh_Hans: 从 Perplexity 获取您的 Perplexity API key + url: https://www.perplexity.ai/settings/api diff --git a/api/core/tools/provider/builtin/perplexity/tools/perplexity_search.py b/api/core/tools/provider/builtin/perplexity/tools/perplexity_search.py new file mode 100644 index 0000000000000000000000000000000000000000..5ed4b9ca9934837981b39098fcd594486f377c8f --- /dev/null +++ b/api/core/tools/provider/builtin/perplexity/tools/perplexity_search.py @@ -0,0 +1,67 @@ +import json +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +PERPLEXITY_API_URL = "https://api.perplexity.ai/chat/completions" + + +class PerplexityAITool(BuiltinTool): + def _parse_response(self, response: dict) -> dict: + """Parse the response from Perplexity AI API""" + if "choices" in response and len(response["choices"]) > 0: + message = response["choices"][0]["message"] + return { + "content": message.get("content", ""), + "role": message.get("role", ""), + "citations": response.get("citations", []), + } + else: + return {"content": "Unable to get a valid response", "role": "assistant", "citations": []} + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + headers = { + "Authorization": f"Bearer {self.runtime.credentials['perplexity_api_key']}", + "Content-Type": "application/json", + } + + payload = { + "model": tool_parameters.get("model", "llama-3.1-sonar-small-128k-online"), + "messages": [ + {"role": "system", "content": "Be precise and concise."}, + {"role": "user", "content": tool_parameters["query"]}, + ], + "max_tokens": tool_parameters.get("max_tokens", 4096), + "temperature": tool_parameters.get("temperature", 0.7), + "top_p": tool_parameters.get("top_p", 1), + "top_k": tool_parameters.get("top_k", 5), + "presence_penalty": tool_parameters.get("presence_penalty", 0), + "frequency_penalty": tool_parameters.get("frequency_penalty", 1), + "stream": False, + } + + if "search_recency_filter" in tool_parameters: + payload["search_recency_filter"] = tool_parameters["search_recency_filter"] + if "return_citations" in tool_parameters: + payload["return_citations"] = tool_parameters["return_citations"] + if "search_domain_filter" in tool_parameters: + if isinstance(tool_parameters["search_domain_filter"], str): + payload["search_domain_filter"] = [tool_parameters["search_domain_filter"]] + elif isinstance(tool_parameters["search_domain_filter"], list): + payload["search_domain_filter"] = tool_parameters["search_domain_filter"] + + response = requests.post(url=PERPLEXITY_API_URL, json=payload, headers=headers) + response.raise_for_status() + valuable_res = self._parse_response(response.json()) + + return [ + self.create_json_message(valuable_res), + self.create_text_message(json.dumps(valuable_res, ensure_ascii=False, indent=2)), + ] diff --git a/api/core/tools/provider/builtin/perplexity/tools/perplexity_search.yaml b/api/core/tools/provider/builtin/perplexity/tools/perplexity_search.yaml new file mode 100644 index 0000000000000000000000000000000000000000..02a645df335aaf51592e67584e5e15160f75925b --- /dev/null +++ b/api/core/tools/provider/builtin/perplexity/tools/perplexity_search.yaml @@ -0,0 +1,178 @@ +identity: + name: perplexity + author: Dify + label: + en_US: Perplexity Search +description: + human: + en_US: Search information using Perplexity AI's language models. + llm: This tool is used to search information using Perplexity AI's language models. +parameters: + - name: query + type: string + required: true + label: + en_US: Query + zh_Hans: 查询 + human_description: + en_US: The text query to be processed by the AI model. + zh_Hans: 要由 AI 模型处理的文本查询。 + form: llm + - name: model + type: select + required: false + label: + en_US: Model Name + zh_Hans: 模型名称 + human_description: + en_US: The Perplexity AI model to use for generating the response. + zh_Hans: 用于生成响应的 Perplexity AI 模型。 + form: form + default: "llama-3.1-sonar-small-128k-online" + options: + - value: llama-3.1-sonar-small-128k-online + label: + en_US: llama-3.1-sonar-small-128k-online + zh_Hans: llama-3.1-sonar-small-128k-online + - value: llama-3.1-sonar-large-128k-online + label: + en_US: llama-3.1-sonar-large-128k-online + zh_Hans: llama-3.1-sonar-large-128k-online + - value: llama-3.1-sonar-huge-128k-online + label: + en_US: llama-3.1-sonar-huge-128k-online + zh_Hans: llama-3.1-sonar-huge-128k-online + - name: max_tokens + type: number + required: false + label: + en_US: Max Tokens + zh_Hans: 最大令牌数 + pt_BR: Máximo de Tokens + human_description: + en_US: The maximum number of tokens to generate in the response. + zh_Hans: 在响应中生成的最大令牌数。 + pt_BR: O número máximo de tokens a serem gerados na resposta. + form: form + default: 4096 + min: 1 + max: 4096 + - name: temperature + type: number + required: false + label: + en_US: Temperature + zh_Hans: 温度 + pt_BR: Temperatura + human_description: + en_US: Controls randomness in the output. Lower values make the output more focused and deterministic. + zh_Hans: 控制输出的随机性。较低的值使输出更加集中和确定。 + form: form + default: 0.7 + min: 0 + max: 1 + - name: top_k + type: number + required: false + label: + en_US: Top K + zh_Hans: 取样数量 + human_description: + en_US: The number of top results to consider for response generation. + zh_Hans: 用于生成响应的顶部结果数量。 + form: form + default: 5 + min: 1 + max: 100 + - name: top_p + type: number + required: false + label: + en_US: Top P + zh_Hans: Top P + human_description: + en_US: Controls diversity via nucleus sampling. + zh_Hans: 通过核心采样控制多样性。 + form: form + default: 1 + min: 0.1 + max: 1 + step: 0.1 + - name: presence_penalty + type: number + required: false + label: + en_US: Presence Penalty + zh_Hans: 存在惩罚 + human_description: + en_US: Positive values penalize new tokens based on whether they appear in the text so far. + zh_Hans: 正值会根据新词元是否已经出现在文本中来对其进行惩罚。 + form: form + default: 0 + min: -1.0 + max: 1.0 + step: 0.1 + - name: frequency_penalty + type: number + required: false + label: + en_US: Frequency Penalty + zh_Hans: 频率惩罚 + human_description: + en_US: Positive values penalize new tokens based on their existing frequency in the text so far. + zh_Hans: 正值会根据新词元在文本中已经出现的频率来对其进行惩罚。 + form: form + default: 1 + min: 0.1 + max: 1.0 + step: 0.1 + - name: return_citations + type: boolean + required: false + label: + en_US: Return Citations + zh_Hans: 返回引用 + human_description: + en_US: Whether to return citations in the response. + zh_Hans: 是否在响应中返回引用。 + form: form + default: true + - name: search_domain_filter + type: string + required: false + label: + en_US: Search Domain Filter + zh_Hans: 搜索域过滤器 + human_description: + en_US: Domain to filter the search results. + zh_Hans: 用于过滤搜索结果的域名。 + form: form + default: "" + - name: search_recency_filter + type: select + required: false + label: + en_US: Search Recency Filter + zh_Hans: 搜索时间过滤器 + human_description: + en_US: Filter for search results based on recency. + zh_Hans: 基于时间筛选搜索结果。 + form: form + default: "month" + options: + - value: day + label: + en_US: Day + zh_Hans: 天 + - value: week + label: + en_US: Week + zh_Hans: 周 + - value: month + label: + en_US: Month + zh_Hans: 月 + - value: year + label: + en_US: Year + zh_Hans: 年 diff --git a/api/core/tools/provider/builtin/podcast_generator/_assets/icon.svg b/api/core/tools/provider/builtin/podcast_generator/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..01743c9cd31120b28e99aa185dd933edf5bc9d37 --- /dev/null +++ b/api/core/tools/provider/builtin/podcast_generator/_assets/icon.svg @@ -0,0 +1,24 @@ + + + + + + + + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/podcast_generator/podcast_generator.py b/api/core/tools/provider/builtin/podcast_generator/podcast_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..a7f7ad2e78b9b0a141cb6de7f832f73c5e6e60ba --- /dev/null +++ b/api/core/tools/provider/builtin/podcast_generator/podcast_generator.py @@ -0,0 +1,38 @@ +from typing import Any + +import openai +from yarl import URL + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class PodcastGeneratorProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + tts_service = credentials.get("tts_service") + api_key = credentials.get("api_key") + base_url = credentials.get("openai_base_url") + + if not tts_service: + raise ToolProviderCredentialValidationError("TTS service is not specified") + + if not api_key: + raise ToolProviderCredentialValidationError("API key is missing") + + if base_url: + base_url = str(URL(base_url) / "v1") + + if tts_service == "openai": + self._validate_openai_credentials(api_key, base_url) + else: + raise ToolProviderCredentialValidationError(f"Unsupported TTS service: {tts_service}") + + def _validate_openai_credentials(self, api_key: str, base_url: str | None) -> None: + client = openai.OpenAI(api_key=api_key, base_url=base_url) + try: + # We're using a simple API call to validate the credentials + client.models.list() + except openai.AuthenticationError: + raise ToolProviderCredentialValidationError("Invalid OpenAI API key") + except Exception as e: + raise ToolProviderCredentialValidationError(f"Error validating OpenAI API key: {str(e)}") diff --git a/api/core/tools/provider/builtin/podcast_generator/podcast_generator.yaml b/api/core/tools/provider/builtin/podcast_generator/podcast_generator.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d4edb17b28638ee39d877a8033d2314aee2b9d5b --- /dev/null +++ b/api/core/tools/provider/builtin/podcast_generator/podcast_generator.yaml @@ -0,0 +1,46 @@ +identity: + author: Dify + name: podcast_generator + label: + en_US: Podcast Generator + zh_Hans: 播客生成器 + description: + en_US: Generate podcast audio using Text-to-Speech services + zh_Hans: 使用文字转语音服务生成播客音频 + icon: icon.svg +credentials_for_provider: + tts_service: + type: select + required: true + label: + en_US: TTS Service + zh_Hans: TTS 服务 + placeholder: + en_US: Select a TTS service + zh_Hans: 选择一个 TTS 服务 + options: + - label: + en_US: OpenAI TTS + zh_Hans: OpenAI TTS + value: openai + api_key: + type: secret-input + required: true + label: + en_US: API Key + zh_Hans: API 密钥 + placeholder: + en_US: Enter your TTS service API key + zh_Hans: 输入您的 TTS 服务 API 密钥 + openai_base_url: + type: text-input + required: false + label: + en_US: OpenAI base URL + zh_Hans: OpenAI base URL + help: + en_US: Please input your OpenAI base URL + zh_Hans: 请输入你的 OpenAI base URL + placeholder: + en_US: Please input your OpenAI base URL + zh_Hans: 请输入你的 OpenAI base URL diff --git a/api/core/tools/provider/builtin/podcast_generator/tools/podcast_audio_generator.py b/api/core/tools/provider/builtin/podcast_generator/tools/podcast_audio_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..704e0015d961a31db3cf543a7b7c7ba893dee4b7 --- /dev/null +++ b/api/core/tools/provider/builtin/podcast_generator/tools/podcast_audio_generator.py @@ -0,0 +1,114 @@ +import concurrent.futures +import io +import random +import warnings +from typing import Any, Literal, Optional, Union + +import openai +from yarl import URL + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.errors import ToolParameterValidationError, ToolProviderCredentialValidationError +from core.tools.tool.builtin_tool import BuiltinTool + +with warnings.catch_warnings(): + warnings.simplefilter("ignore") + from pydub import AudioSegment # type: ignore + + +class PodcastAudioGeneratorTool(BuiltinTool): + @staticmethod + def _generate_silence(duration: float): + # Generate silent WAV data using pydub + silence = AudioSegment.silent(duration=int(duration * 1000)) # pydub uses milliseconds + return silence + + @staticmethod + def _generate_audio_segment( + client: openai.OpenAI, + line: str, + voice: Literal["alloy", "echo", "fable", "onyx", "nova", "shimmer"], + index: int, + ) -> tuple[int, Union[AudioSegment, str], Optional[AudioSegment]]: + try: + response = client.audio.speech.create(model="tts-1", voice=voice, input=line.strip(), response_format="wav") + audio = AudioSegment.from_wav(io.BytesIO(response.content)) + silence_duration = random.uniform(0.1, 1.5) + silence = PodcastAudioGeneratorTool._generate_silence(silence_duration) + return index, audio, silence + except Exception as e: + return index, f"Error generating audio: {str(e)}", None + + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + # Extract parameters + script = tool_parameters.get("script", "") + host1_voice = tool_parameters.get("host1_voice") + host2_voice = tool_parameters.get("host2_voice") + + # Split the script into lines + script_lines = [line for line in script.split("\n") if line.strip()] + + # Ensure voices are provided + if not host1_voice or not host2_voice: + raise ToolParameterValidationError("Host voices are required") + + # Ensure runtime and credentials + if not self.runtime or not self.runtime.credentials: + raise ToolProviderCredentialValidationError("Tool runtime or credentials are missing") + + # Get OpenAI API key from credentials + api_key = self.runtime.credentials.get("api_key") + if not api_key: + raise ToolProviderCredentialValidationError("OpenAI API key is missing") + + # Get OpenAI base URL + openai_base_url = self.runtime.credentials.get("openai_base_url", None) + openai_base_url = str(URL(openai_base_url) / "v1") if openai_base_url else None + + # Initialize OpenAI client + client = openai.OpenAI( + api_key=api_key, + base_url=openai_base_url, + ) + + # Create a thread pool + max_workers = 5 + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [] + for i, line in enumerate(script_lines): + voice = host1_voice if i % 2 == 0 else host2_voice + future = executor.submit(self._generate_audio_segment, client, line, voice, i) + futures.append(future) + + # Collect results + audio_segments: list[Any] = [None] * len(script_lines) + for future in concurrent.futures.as_completed(futures): + index, audio, silence = future.result() + if isinstance(audio, str): # Error occurred + return self.create_text_message(audio) + audio_segments[index] = (audio, silence) + + # Combine audio segments in the correct order + combined_audio = AudioSegment.empty() + for i, (audio, silence) in enumerate(audio_segments): + if audio: + combined_audio += audio + if i < len(audio_segments) - 1 and silence: + combined_audio += silence + + # Export the combined audio to a WAV file in memory + buffer = io.BytesIO() + combined_audio.export(buffer, format="wav") + wav_bytes = buffer.getvalue() + + # Create a blob message with the combined audio + return [ + self.create_text_message("Audio generated successfully"), + self.create_blob_message( + blob=wav_bytes, + meta={"mime_type": "audio/x-wav"}, + save_as=self.VariableKey.AUDIO, + ), + ] diff --git a/api/core/tools/provider/builtin/podcast_generator/tools/podcast_audio_generator.yaml b/api/core/tools/provider/builtin/podcast_generator/tools/podcast_audio_generator.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d6ae98f59522c580a7363bb43d22cbb9c5ab5021 --- /dev/null +++ b/api/core/tools/provider/builtin/podcast_generator/tools/podcast_audio_generator.yaml @@ -0,0 +1,95 @@ +identity: + name: podcast_audio_generator + author: Dify + label: + en_US: Podcast Audio Generator + zh_Hans: 播客音频生成器 +description: + human: + en_US: Generate a podcast audio file from a script with two alternating voices using OpenAI's TTS service. + zh_Hans: 使用 OpenAI 的 TTS 服务,从包含两个交替声音的脚本生成播客音频文件。 + llm: This tool converts a prepared podcast script into an audio file using OpenAI's Text-to-Speech service, with two specified voices for alternating hosts. +parameters: + - name: script + type: string + required: true + label: + en_US: Podcast Script + zh_Hans: 播客脚本 + human_description: + en_US: A string containing alternating lines for two hosts, separated by newline characters. + zh_Hans: 包含两位主持人交替台词的字符串,每行用换行符分隔。 + llm_description: A string representing the script, with alternating lines for two hosts separated by newline characters. + form: llm + - name: host1_voice + type: select + required: true + label: + en_US: Host 1 Voice + zh_Hans: 主持人1 音色 + human_description: + en_US: The voice for the first host. + zh_Hans: 第一位主持人的音色。 + llm_description: The voice identifier for the first host's voice. + options: + - label: + en_US: Alloy + zh_Hans: Alloy + value: alloy + - label: + en_US: Echo + zh_Hans: Echo + value: echo + - label: + en_US: Fable + zh_Hans: Fable + value: fable + - label: + en_US: Onyx + zh_Hans: Onyx + value: onyx + - label: + en_US: Nova + zh_Hans: Nova + value: nova + - label: + en_US: Shimmer + zh_Hans: Shimmer + value: shimmer + form: form + - name: host2_voice + type: select + required: true + label: + en_US: Host 2 Voice + zh_Hans: 主持人2 音色 + human_description: + en_US: The voice for the second host. + zh_Hans: 第二位主持人的音色。 + llm_description: The voice identifier for the second host's voice. + options: + - label: + en_US: Alloy + zh_Hans: Alloy + value: alloy + - label: + en_US: Echo + zh_Hans: Echo + value: echo + - label: + en_US: Fable + zh_Hans: Fable + value: fable + - label: + en_US: Onyx + zh_Hans: Onyx + value: onyx + - label: + en_US: Nova + zh_Hans: Nova + value: nova + - label: + en_US: Shimmer + zh_Hans: Shimmer + value: shimmer + form: form diff --git a/api/core/tools/provider/builtin/pubmed/_assets/icon.svg b/api/core/tools/provider/builtin/pubmed/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..6d6ff593f0c9991fb18289bb23e152e5ff45576e --- /dev/null +++ b/api/core/tools/provider/builtin/pubmed/_assets/icon.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/pubmed/pubmed.py b/api/core/tools/provider/builtin/pubmed/pubmed.py new file mode 100644 index 0000000000000000000000000000000000000000..ea3a477c30178d5779ba0aa397703e3f7b6f15dd --- /dev/null +++ b/api/core/tools/provider/builtin/pubmed/pubmed.py @@ -0,0 +1,20 @@ +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.pubmed.tools.pubmed_search import PubMedSearchTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class PubMedProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + try: + PubMedSearchTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "query": "John Doe", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/pubmed/pubmed.yaml b/api/core/tools/provider/builtin/pubmed/pubmed.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5f8303147c397be17cc0ab30c2d2f801dfd89b7e --- /dev/null +++ b/api/core/tools/provider/builtin/pubmed/pubmed.yaml @@ -0,0 +1,13 @@ +identity: + author: Pink Banana + name: pubmed + label: + en_US: PubMed + zh_Hans: PubMed + description: + en_US: A search engine for biomedical literature. + zh_Hans: 一款生物医学文献搜索引擎。 + icon: icon.svg + tags: + - medical + - search diff --git a/api/core/tools/provider/builtin/pubmed/tools/pubmed_search.py b/api/core/tools/provider/builtin/pubmed/tools/pubmed_search.py new file mode 100644 index 0000000000000000000000000000000000000000..3a4f374ea0b0bca66e889e021b50f92c348754eb --- /dev/null +++ b/api/core/tools/provider/builtin/pubmed/tools/pubmed_search.py @@ -0,0 +1,191 @@ +import json +import time +import urllib.error +import urllib.parse +import urllib.request +from typing import Any + +from pydantic import BaseModel, Field + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class PubMedAPIWrapper(BaseModel): + """ + Wrapper around PubMed API. + + This wrapper will use the PubMed API to conduct searches and fetch + document summaries. By default, it will return the document summaries + of the top-k results of an input search. + + Parameters: + top_k_results: number of the top-scored document used for the PubMed tool + load_max_docs: a limit to the number of loaded documents + load_all_available_meta: + if True: the `metadata` of the loaded Documents gets all available meta info + (see https://www.ncbi.nlm.nih.gov/books/NBK25499/#chapter4.ESearch) + if False: the `metadata` gets only the most informative fields. + """ + + base_url_esearch: str = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi?" + base_url_efetch: str = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi?" + max_retry: int = 5 + sleep_time: float = 0.2 + + # Default values for the parameters + top_k_results: int = 3 + load_max_docs: int = 25 + ARXIV_MAX_QUERY_LENGTH: int = 300 + doc_content_chars_max: int = 2000 + load_all_available_meta: bool = False + email: str = "your_email@example.com" + + def run(self, query: str) -> str: + """ + Run PubMed search and get the article meta information. + See https://www.ncbi.nlm.nih.gov/books/NBK25499/#chapter4.ESearch + It uses only the most informative fields of article meta information. + """ + + try: + # Retrieve the top-k results for the query + docs = [ + f"Published: {result['pub_date']}\nTitle: {result['title']}\nSummary: {result['summary']}" + for result in self.load(query[: self.ARXIV_MAX_QUERY_LENGTH]) + ] + + # Join the results and limit the character count + return "\n\n".join(docs)[: self.doc_content_chars_max] if docs else "No good PubMed Result was found" + except Exception as ex: + return f"PubMed exception: {ex}" + + def load(self, query: str) -> list[dict]: + """ + Search PubMed for documents matching the query. + Return a list of dictionaries containing the document metadata. + """ + + url = ( + self.base_url_esearch + + "db=pubmed&term=" + + str({urllib.parse.quote(query)}) + + f"&retmode=json&retmax={self.top_k_results}&usehistory=y" + ) + result = urllib.request.urlopen(url) + text = result.read().decode("utf-8") + json_text = json.loads(text) + + articles = [] + webenv = json_text["esearchresult"]["webenv"] + for uid in json_text["esearchresult"]["idlist"]: + article = self.retrieve_article(uid, webenv) + articles.append(article) + + # Convert the list of articles to a JSON string + return articles + + def retrieve_article(self, uid: str, webenv: str) -> dict: + url = self.base_url_efetch + "db=pubmed&retmode=xml&id=" + uid + "&webenv=" + webenv + + retry = 0 + while True: + try: + result = urllib.request.urlopen(url) + break + except urllib.error.HTTPError as e: + if e.code == 429 and retry < self.max_retry: + # Too Many Requests error + # wait for an exponentially increasing amount of time + print(f"Too Many Requests, waiting for {self.sleep_time:.2f} seconds...") + time.sleep(self.sleep_time) + self.sleep_time *= 2 + retry += 1 + else: + raise e + + xml_text = result.read().decode("utf-8") + + # Get title + title = "" + if "" in xml_text and "" in xml_text: + start_tag = "" + end_tag = "" + title = xml_text[xml_text.index(start_tag) + len(start_tag) : xml_text.index(end_tag)] + + # Get abstract + abstract = "" + if "" in xml_text and "" in xml_text: + start_tag = "" + end_tag = "" + abstract = xml_text[xml_text.index(start_tag) + len(start_tag) : xml_text.index(end_tag)] + + # Get publication date + pub_date = "" + if "" in xml_text and "" in xml_text: + start_tag = "" + end_tag = "" + pub_date = xml_text[xml_text.index(start_tag) + len(start_tag) : xml_text.index(end_tag)] + + # Return article as dictionary + article = { + "uid": uid, + "title": title, + "summary": abstract, + "pub_date": pub_date, + } + return article + + +class PubmedQueryRun(BaseModel): + """Tool that searches the PubMed API.""" + + name: str = "PubMed" + description: str = ( + "A wrapper around PubMed.org " + "Useful for when you need to answer questions about Physics, Mathematics, " + "Computer Science, Quantitative Biology, Quantitative Finance, Statistics, " + "Electrical Engineering, and Economics " + "from scientific articles on PubMed.org. " + "Input should be a search query." + ) + api_wrapper: PubMedAPIWrapper = Field(default_factory=PubMedAPIWrapper) + + def _run( + self, + query: str, + ) -> str: + """Use the Arxiv tool.""" + return self.api_wrapper.run(query) + + +class PubMedInput(BaseModel): + query: str = Field(..., description="Search query.") + + +class PubMedSearchTool(BuiltinTool): + """ + Tool for performing a search using PubMed search engine. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage | list[ToolInvokeMessage]: + """ + Invoke the PubMed search tool. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (dict[str, Any]): The parameters for the tool invocation. + + Returns: + ToolInvokeMessage | list[ToolInvokeMessage]: The result of the tool invocation. + """ + query = tool_parameters.get("query", "") + + if not query: + return self.create_text_message("Please input query") + + tool = PubmedQueryRun(args_schema=PubMedInput) + + result = tool._run(query) + + return self.create_text_message(self.summary(user_id=user_id, content=result)) diff --git a/api/core/tools/provider/builtin/pubmed/tools/pubmed_search.yaml b/api/core/tools/provider/builtin/pubmed/tools/pubmed_search.yaml new file mode 100644 index 0000000000000000000000000000000000000000..77ab809fbc3e051c194471d708fd4a7670b47adf --- /dev/null +++ b/api/core/tools/provider/builtin/pubmed/tools/pubmed_search.yaml @@ -0,0 +1,23 @@ +identity: + name: pubmed_search + author: Pink Banana + label: + en_US: PubMed Search + zh_Hans: PubMed 搜索 +description: + human: + en_US: PubMed® comprises more than 35 million citations for biomedical literature from MEDLINE, life science journals, and online books. Citations may include links to full text content from PubMed Central and publisher web sites. + zh_Hans: PubMed® 包含来自 MEDLINE、生命科学期刊和在线书籍的超过 3500 万篇生物医学文献引用。引用可能包括来自 PubMed Central 和出版商网站的全文内容链接。 + llm: Perform searches on PubMed and get results. +parameters: + - name: query + type: string + required: true + label: + en_US: Query string + zh_Hans: 查询语句 + human_description: + en_US: The search query. + zh_Hans: 搜索查询语句。 + llm_description: Key words for searching + form: llm diff --git a/api/core/tools/provider/builtin/qrcode/_assets/icon.svg b/api/core/tools/provider/builtin/qrcode/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..979bdda45582561907f287552110e311beb31e72 --- /dev/null +++ b/api/core/tools/provider/builtin/qrcode/_assets/icon.svg @@ -0,0 +1,7 @@ + + + + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/qrcode/qrcode.py b/api/core/tools/provider/builtin/qrcode/qrcode.py new file mode 100644 index 0000000000000000000000000000000000000000..8466b9a26b42b6ad63009431e02e8c1f9da4ada1 --- /dev/null +++ b/api/core/tools/provider/builtin/qrcode/qrcode.py @@ -0,0 +1,13 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.qrcode.tools.qrcode_generator import QRCodeGeneratorTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class QRCodeProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + QRCodeGeneratorTool().invoke(user_id="", tool_parameters={"content": "Dify 123 😊"}) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/qrcode/qrcode.yaml b/api/core/tools/provider/builtin/qrcode/qrcode.yaml new file mode 100644 index 0000000000000000000000000000000000000000..82e2a06e15cc18db5ace4591fc989e02c8b29046 --- /dev/null +++ b/api/core/tools/provider/builtin/qrcode/qrcode.yaml @@ -0,0 +1,14 @@ +identity: + author: Bowen Liang + name: qrcode + label: + en_US: QRCode + zh_Hans: 二维码工具 + pt_BR: QRCode + description: + en_US: A tool for generating QR code (quick-response code) image. + zh_Hans: 一个二维码工具 + pt_BR: A tool for generating QR code (quick-response code) image. + icon: icon.svg + tags: + - utilities diff --git a/api/core/tools/provider/builtin/qrcode/tools/qrcode_generator.py b/api/core/tools/provider/builtin/qrcode/tools/qrcode_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..4a47c4211f4fd420925bbbe6c40be61857cf8ee5 --- /dev/null +++ b/api/core/tools/provider/builtin/qrcode/tools/qrcode_generator.py @@ -0,0 +1,70 @@ +import io +import logging +from typing import Any, Union + +from qrcode.constants import ERROR_CORRECT_H, ERROR_CORRECT_L, ERROR_CORRECT_M, ERROR_CORRECT_Q # type: ignore +from qrcode.image.base import BaseImage # type: ignore +from qrcode.image.pure import PyPNGImage # type: ignore +from qrcode.main import QRCode # type: ignore + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class QRCodeGeneratorTool(BuiltinTool): + error_correction_levels: dict[str, int] = { + "L": ERROR_CORRECT_L, # <=7% + "M": ERROR_CORRECT_M, # <=15% + "Q": ERROR_CORRECT_Q, # <=25% + "H": ERROR_CORRECT_H, # <=30% + } + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + # get text content + content = tool_parameters.get("content", "") + if not content: + return self.create_text_message("Invalid parameter content") + + # get border size + border = tool_parameters.get("border", 0) + if border < 0 or border > 100: + return self.create_text_message("Invalid parameter border") + + # get error_correction + error_correction = tool_parameters.get("error_correction", "") + if error_correction not in self.error_correction_levels: + return self.create_text_message("Invalid parameter error_correction") + + try: + image = self._generate_qrcode(content, border, error_correction) + image_bytes = self._image_to_byte_array(image) + return self.create_blob_message( + blob=image_bytes, meta={"mime_type": "image/png"}, save_as=self.VariableKey.IMAGE.value + ) + except Exception: + logging.exception(f"Failed to generate QR code for content: {content}") + return self.create_text_message("Failed to generate QR code") + + def _generate_qrcode(self, content: str, border: int, error_correction: str) -> BaseImage: + qr = QRCode( + image_factory=PyPNGImage, + error_correction=self.error_correction_levels.get(error_correction), + border=border, + ) + qr.add_data(data=content) + qr.make(fit=True) + img = qr.make_image() + return img + + @staticmethod + def _image_to_byte_array(image: BaseImage) -> bytes: + byte_stream = io.BytesIO() + image.save(byte_stream) + return byte_stream.getvalue() diff --git a/api/core/tools/provider/builtin/qrcode/tools/qrcode_generator.yaml b/api/core/tools/provider/builtin/qrcode/tools/qrcode_generator.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8c8b8c449ad5135faac9f4c5820ca5575fc99c9f --- /dev/null +++ b/api/core/tools/provider/builtin/qrcode/tools/qrcode_generator.yaml @@ -0,0 +1,76 @@ +identity: + name: qrcode_generator + author: Bowen Liang + label: + en_US: Generate QR Code + zh_Hans: 生成二维码 + pt_BR: Generate QR Code +description: + human: + en_US: A tool for generating QR code image + zh_Hans: 一个用于生成二维码的工具 + pt_BR: A tool for generating QR code image + llm: A tool for generating QR code image +parameters: + - name: content + type: string + required: true + label: + en_US: content text for QR code + zh_Hans: 二维码文本内容 + pt_BR: content text for QR code + human_description: + en_US: content text for QR code + zh_Hans: 二维码文本内容 + pt_BR: 二维码文本内容 + form: llm + - name: error_correction + type: select + required: true + default: M + label: + en_US: Error Correction + zh_Hans: 容错等级 + pt_BR: Error Correction + human_description: + en_US: Error Correction in L, M, Q or H, from low to high, the bigger size of generated QR code with the better error correction effect + zh_Hans: 容错等级,可设置为低、中、偏高或高,从低到高,生成的二维码越大且容错效果越好 + pt_BR: Error Correction in L, M, Q or H, from low to high, the bigger size of generated QR code with the better error correction effect + options: + - value: L + label: + en_US: Low + zh_Hans: 低 + pt_BR: Low + - value: M + label: + en_US: Medium + zh_Hans: 中 + pt_BR: Medium + - value: Q + label: + en_US: Quartile + zh_Hans: 偏高 + pt_BR: Quartile + - value: H + label: + en_US: High + zh_Hans: 高 + pt_BR: High + form: form + - name: border + type: number + required: true + default: 2 + min: 0 + max: 100 + label: + en_US: border size + zh_Hans: 边框粗细 + pt_BR: border size + human_description: + en_US: border size(default to 2) + zh_Hans: 边框粗细的格数(默认为2) + pt_BR: border size(default to 2) + llm: border size, default to 2 + form: form diff --git a/api/core/tools/provider/builtin/rapidapi/_assets/rapidapi.png b/api/core/tools/provider/builtin/rapidapi/_assets/rapidapi.png new file mode 100644 index 0000000000000000000000000000000000000000..9c7468bb172326d7649a2f911c6f5dd7aa75584f Binary files /dev/null and b/api/core/tools/provider/builtin/rapidapi/_assets/rapidapi.png differ diff --git a/api/core/tools/provider/builtin/rapidapi/rapidapi.py b/api/core/tools/provider/builtin/rapidapi/rapidapi.py new file mode 100644 index 0000000000000000000000000000000000000000..31077b0894153b4b575891a9dec2510ecfc2e0ed --- /dev/null +++ b/api/core/tools/provider/builtin/rapidapi/rapidapi.py @@ -0,0 +1,22 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.rapidapi.tools.google_news import GooglenewsTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class RapidapiProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + GooglenewsTool().fork_tool_runtime( + meta={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "language_region": "en-US", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/rapidapi/rapidapi.yaml b/api/core/tools/provider/builtin/rapidapi/rapidapi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3f1d1c58248cc1e0df3444667f1c6771f09a0e2b --- /dev/null +++ b/api/core/tools/provider/builtin/rapidapi/rapidapi.yaml @@ -0,0 +1,39 @@ +identity: + name: rapidapi + author: Steven Sun + label: + en_US: RapidAPI + zh_Hans: RapidAPI + description: + en_US: RapidAPI is the world's largest API marketplace with over 1,000,000 developers and 10,000 APIs. + zh_Hans: RapidAPI是全球最大的API市场,拥有超过100万开发人员和10000个API。 + icon: rapidapi.png + tags: + - news +credentials_for_provider: + x-rapidapi-host: + type: text-input + required: true + label: + en_US: x-rapidapi-host + zh_Hans: x-rapidapi-host + placeholder: + en_US: Please input your x-rapidapi-host + zh_Hans: 请输入你的 x-rapidapi-host + help: + en_US: Get your x-rapidapi-host from RapidAPI. + zh_Hans: 从 RapidAPI 获取您的 x-rapidapi-host。 + url: https://rapidapi.com/ + x-rapidapi-key: + type: secret-input + required: true + label: + en_US: x-rapidapi-key + zh_Hans: x-rapidapi-key + placeholder: + en_US: Please input your x-rapidapi-key + zh_Hans: 请输入你的 x-rapidapi-key + help: + en_US: Get your x-rapidapi-key from RapidAPI. + zh_Hans: 从 RapidAPI 获取您的 x-rapidapi-key。 + url: https://rapidapi.com/ diff --git a/api/core/tools/provider/builtin/rapidapi/tools/google_news.py b/api/core/tools/provider/builtin/rapidapi/tools/google_news.py new file mode 100644 index 0000000000000000000000000000000000000000..d4b6dc4a46b6ef357901624b25dfb0b3e6ec6ea5 --- /dev/null +++ b/api/core/tools/provider/builtin/rapidapi/tools/google_news.py @@ -0,0 +1,33 @@ +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.errors import ToolInvokeError, ToolProviderCredentialValidationError +from core.tools.tool.builtin_tool import BuiltinTool + + +class GooglenewsTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + key = self.runtime.credentials.get("x-rapidapi-key", "") + host = self.runtime.credentials.get("x-rapidapi-host", "") + if not all([key, host]): + raise ToolProviderCredentialValidationError("Please input correct x-rapidapi-key and x-rapidapi-host") + headers = {"x-rapidapi-key": key, "x-rapidapi-host": host} + lr = tool_parameters.get("language_region", "") + url = f"https://{host}/latest?lr={lr}" + response = requests.get(url, headers=headers) + if response.status_code != 200: + raise ToolInvokeError(f"Error {response.status_code}: {response.text}") + return self.create_text_message(response.text) + + def validate_credentials(self, parameters: dict[str, Any]) -> None: + parameters["validate"] = True + self._invoke(parameters) diff --git a/api/core/tools/provider/builtin/rapidapi/tools/google_news.yaml b/api/core/tools/provider/builtin/rapidapi/tools/google_news.yaml new file mode 100644 index 0000000000000000000000000000000000000000..547681b16663d55ce1f771a6856dfaae69d2b605 --- /dev/null +++ b/api/core/tools/provider/builtin/rapidapi/tools/google_news.yaml @@ -0,0 +1,24 @@ +identity: + name: google_news + author: Steven Sun + label: + en_US: GoogleNews + zh_Hans: 谷歌新闻 +description: + human: + en_US: google news is a news aggregator service developed by Google. It presents a continuous, customizable flow of articles organized from thousands of publishers and magazines. + zh_Hans: 谷歌新闻是由谷歌开发的新闻聚合服务。它提供了一个持续的、可定制的文章流,这些文章是从成千上万的出版商和杂志中整理出来的。 + llm: A tool to get the latest news from Google News. +parameters: + - name: language_region + type: string + required: true + label: + en_US: Language and Region + zh_Hans: 语言和地区 + human_description: + en_US: The language and region determine the language and region of the search results, and its value is assigned according to the "National Language Code Comparison Table", such as en-US, which stands for English (United States); zh-CN, stands for Chinese (Simplified). + zh_Hans: 语言和地区决定了搜索结果的语言和地区,其赋值按照《国家语言代码对照表》,形如en-US,代表英语(美国);zh-CN,代表中文(简体)。 + llm_description: The language and region determine the language and region of the search results, and its value is assigned according to the "National Language Code Comparison Table", such as en-US, which stands for English (United States); zh-CN, stands for Chinese (Simplified). + default: en-US + form: llm diff --git a/api/core/tools/provider/builtin/regex/_assets/icon.svg b/api/core/tools/provider/builtin/regex/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..0231a2b4aa9da24cf58cf0054bd92c2e1499c41c --- /dev/null +++ b/api/core/tools/provider/builtin/regex/_assets/icon.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/regex/regex.py b/api/core/tools/provider/builtin/regex/regex.py new file mode 100644 index 0000000000000000000000000000000000000000..c498105979f13eb9dfc09df753e465b377281e2d --- /dev/null +++ b/api/core/tools/provider/builtin/regex/regex.py @@ -0,0 +1,19 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.regex.tools.regex_extract import RegexExpressionTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class RegexProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + RegexExpressionTool().invoke( + user_id="", + tool_parameters={ + "content": "1+(2+3)*4", + "expression": r"(\d+)", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/regex/regex.yaml b/api/core/tools/provider/builtin/regex/regex.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d05776f214e8d2ffff3b2944948236ba55055e78 --- /dev/null +++ b/api/core/tools/provider/builtin/regex/regex.yaml @@ -0,0 +1,15 @@ +identity: + author: zhuhao + name: regex + label: + en_US: Regex + zh_Hans: 正则表达式提取 + pt_BR: Regex + description: + en_US: A tool for regex extraction. + zh_Hans: 一个用于正则表达式内容提取的工具。 + pt_BR: A tool for regex extraction. + icon: icon.svg + tags: + - utilities + - productivity diff --git a/api/core/tools/provider/builtin/regex/tools/regex_extract.py b/api/core/tools/provider/builtin/regex/tools/regex_extract.py new file mode 100644 index 0000000000000000000000000000000000000000..786b46940400300889b2f895b4a717fc529e4aa6 --- /dev/null +++ b/api/core/tools/provider/builtin/regex/tools/regex_extract.py @@ -0,0 +1,28 @@ +import re +from typing import Any, Union + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class RegexExpressionTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + # get expression + content = tool_parameters.get("content", "").strip() + if not content: + return self.create_text_message("Invalid content") + expression = tool_parameters.get("expression", "").strip() + if not expression: + return self.create_text_message("Invalid expression") + try: + result = re.findall(expression, content) + return self.create_text_message(str(result)) + except Exception as e: + return self.create_text_message(f"Failed to extract result, error: {str(e)}") diff --git a/api/core/tools/provider/builtin/regex/tools/regex_extract.yaml b/api/core/tools/provider/builtin/regex/tools/regex_extract.yaml new file mode 100644 index 0000000000000000000000000000000000000000..de4100def176c9408354e38e09b4f44cadce4f8e --- /dev/null +++ b/api/core/tools/provider/builtin/regex/tools/regex_extract.yaml @@ -0,0 +1,38 @@ +identity: + name: regex_extract + author: zhuhao + label: + en_US: Regex Extract + zh_Hans: 正则表达式内容提取 + pt_BR: Regex Extract +description: + human: + en_US: A tool for extracting matching content using regular expressions. + zh_Hans: 一个用于利用正则表达式提取匹配内容结果的工具。 + pt_BR: A tool for extracting matching content using regular expressions. + llm: A tool for extracting matching content using regular expressions. +parameters: + - name: content + type: string + required: true + label: + en_US: Content to be extracted + zh_Hans: 内容 + pt_BR: Content to be extracted + human_description: + en_US: Content to be extracted + zh_Hans: 内容 + pt_BR: Content to be extracted + form: llm + - name: expression + type: string + required: true + label: + en_US: Regular expression + zh_Hans: 正则表达式 + pt_BR: Regular expression + human_description: + en_US: Regular expression + zh_Hans: 正则表达式 + pt_BR: Regular expression + form: llm diff --git a/api/core/tools/provider/builtin/searchapi/_assets/icon.svg b/api/core/tools/provider/builtin/searchapi/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..7660b2f351c43bbc1b4e33c15ad12b8128d1f81b --- /dev/null +++ b/api/core/tools/provider/builtin/searchapi/_assets/icon.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/searchapi/searchapi.py b/api/core/tools/provider/builtin/searchapi/searchapi.py new file mode 100644 index 0000000000000000000000000000000000000000..109bba8b2d8f7918a2eb3df9c2d75c221566cf44 --- /dev/null +++ b/api/core/tools/provider/builtin/searchapi/searchapi.py @@ -0,0 +1,20 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.searchapi.tools.google import GoogleTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class SearchAPIProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + GoogleTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={"query": "SearchApi dify", "result_type": "link"}, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/searchapi/searchapi.yaml b/api/core/tools/provider/builtin/searchapi/searchapi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c2fa3f398e192fc5e33979192e323d2dfccb8cc1 --- /dev/null +++ b/api/core/tools/provider/builtin/searchapi/searchapi.yaml @@ -0,0 +1,34 @@ +identity: + author: SearchApi + name: searchapi + label: + en_US: SearchApi + zh_Hans: SearchApi + pt_BR: SearchApi + description: + en_US: SearchApi is a robust real-time SERP API delivering structured data from a collection of search engines including Google Search, Google Jobs, YouTube, Google News, and many more. + zh_Hans: SearchApi 是一个强大的实时 SERP API,可提供来自 Google 搜索、Google 招聘、YouTube、Google 新闻等搜索引擎集合的结构化数据。 + pt_BR: SearchApi is a robust real-time SERP API delivering structured data from a collection of search engines including Google Search, Google Jobs, YouTube, Google News, and many more. + icon: icon.svg + tags: + - search + - business + - news + - productivity +credentials_for_provider: + searchapi_api_key: + type: secret-input + required: true + label: + en_US: SearchApi API key + zh_Hans: SearchApi API key + pt_BR: SearchApi API key + placeholder: + en_US: Please input your SearchApi API key + zh_Hans: 请输入你的 SearchApi API key + pt_BR: Please input your SearchApi API key + help: + en_US: Get your SearchApi API key from SearchApi + zh_Hans: 从 SearchApi 获取您的 SearchApi API key + pt_BR: Get your SearchApi API key from SearchApi + url: https://www.searchapi.io/ diff --git a/api/core/tools/provider/builtin/searchapi/tools/google.py b/api/core/tools/provider/builtin/searchapi/tools/google.py new file mode 100644 index 0000000000000000000000000000000000000000..29d36f5f2326946cf30315b39a96a88867d6d0cf --- /dev/null +++ b/api/core/tools/provider/builtin/searchapi/tools/google.py @@ -0,0 +1,112 @@ +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +SEARCH_API_URL = "https://www.searchapi.io/api/v1/search" + + +class SearchAPI: + """ + SearchAPI tool provider. + """ + + def __init__(self, api_key: str) -> None: + """Initialize SearchAPI tool provider.""" + self.searchapi_api_key = api_key + + def run(self, query: str, **kwargs: Any) -> str: + """Run query through SearchAPI and parse result.""" + type = kwargs.get("result_type", "text") + return self._process_response(self.results(query, **kwargs), type=type) + + def results(self, query: str, **kwargs: Any) -> dict: + """Run query through SearchAPI and return the raw result.""" + params = self.get_params(query, **kwargs) + response = requests.get( + url=SEARCH_API_URL, + params=params, + headers={"Authorization": f"Bearer {self.searchapi_api_key}"}, + ) + response.raise_for_status() + return response.json() + + def get_params(self, query: str, **kwargs: Any) -> dict[str, str]: + """Get parameters for SearchAPI.""" + return { + "engine": "google", + "q": query, + **{key: value for key, value in kwargs.items() if value not in {None, ""}}, + } + + @staticmethod + def _process_response(res: dict, type: str) -> str: + """Process response from SearchAPI.""" + if "error" in res: + return res["error"] + + toret = "" + if type == "text": + if "answer_box" in res and "answer" in res["answer_box"]: + toret += res["answer_box"]["answer"] + "\n" + if "answer_box" in res and "snippet" in res["answer_box"]: + toret += res["answer_box"]["snippet"] + "\n" + if "knowledge_graph" in res and "description" in res["knowledge_graph"]: + toret += res["knowledge_graph"]["description"] + "\n" + if "organic_results" in res and "snippet" in res["organic_results"][0]: + for item in res["organic_results"]: + toret += "content: " + item["snippet"] + "\n" + "link: " + item["link"] + "\n" + if toret == "": + toret = "No good search result found" + + elif type == "link": + if "answer_box" in res and "organic_result" in res["answer_box"]: + if "title" in res["answer_box"]["organic_result"]: + toret = ( + f"[{res['answer_box']['organic_result']['title']}]" + f"({res['answer_box']['organic_result']['link']})\n" + ) + elif "organic_results" in res and "link" in res["organic_results"][0]: + toret = "" + for item in res["organic_results"]: + toret += f"[{item['title']}]({item['link']})\n" + elif "related_questions" in res and "link" in res["related_questions"][0]: + toret = "" + for item in res["related_questions"]: + toret += f"[{item['title']}]({item['link']})\n" + elif "related_searches" in res and "link" in res["related_searches"][0]: + toret = "" + for item in res["related_searches"]: + toret += f"[{item['title']}]({item['link']})\n" + else: + toret = "No good search result found" + return toret + + +class GoogleTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Invoke the SearchApi tool. + """ + query = tool_parameters["query"] + result_type = tool_parameters["result_type"] + num = tool_parameters.get("num", 10) + google_domain = tool_parameters.get("google_domain", "google.com") + gl = tool_parameters.get("gl", "us") + hl = tool_parameters.get("hl", "en") + location = tool_parameters.get("location") + + api_key = self.runtime.credentials["searchapi_api_key"] + result = SearchAPI(api_key).run( + query, result_type=result_type, num=num, google_domain=google_domain, gl=gl, hl=hl, location=location + ) + + if result_type == "text": + return self.create_text_message(text=result) + return self.create_link_message(link=result) diff --git a/api/core/tools/provider/builtin/searchapi/tools/google.yaml b/api/core/tools/provider/builtin/searchapi/tools/google.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0dc1b6672436cdb7294028ae25193b7c3e9afeb9 --- /dev/null +++ b/api/core/tools/provider/builtin/searchapi/tools/google.yaml @@ -0,0 +1,1921 @@ +identity: + name: google_search_api + author: SearchApi + label: + en_US: Google Search API + zh_Hans: Google Search API +description: + human: + en_US: A tool to retrieve answer boxes, knowledge graphs, snippets, and webpages from Google Search engine. + zh_Hans: 一种从 Google 搜索引擎检索答案框、知识图、片段和网页的工具。 + llm: A tool to retrieve answer boxes, knowledge graphs, snippets, and webpages from Google Search engine. +parameters: + - name: query + type: string + required: true + label: + en_US: Query + zh_Hans: 询问 + human_description: + en_US: Defines the query you want to search. + zh_Hans: 定义您要搜索的查询。 + llm_description: Defines the search query you want to search. + form: llm + - name: result_type + type: select + required: true + options: + - value: text + label: + en_US: text + zh_Hans: 文本 + - value: link + label: + en_US: link + zh_Hans: 链接 + default: text + label: + en_US: Result type + zh_Hans: 结果类型 + human_description: + en_US: used for selecting the result type, text or link + zh_Hans: 用于选择结果类型,使用文本还是链接进行展示 + form: form + - name: location + type: string + required: false + label: + en_US: Location + zh_Hans: 询问 + human_description: + en_US: Defines from where you want the search to originate. (For example - New York) + zh_Hans: 定义您想要搜索的起始位置。 (例如 - 纽约) + llm_description: Defines from where you want the search to originate. (For example - New York) + form: llm + - name: gl + type: select + label: + en_US: Country + zh_Hans: 国家 + required: false + human_description: + en_US: Defines the country of the search. Default is "US". + zh_Hans: 定义搜索的国家/地区。默认为“美国”。 + llm_description: Defines the gl parameter of the Google search. + form: form + default: US + options: + - value: AF + label: + en_US: Afghanistan + zh_Hans: 阿富汗 + pt_BR: Afeganistão + - value: AL + label: + en_US: Albania + zh_Hans: 阿尔巴尼亚 + pt_BR: Albânia + - value: DZ + label: + en_US: Algeria + zh_Hans: 阿尔及利亚 + pt_BR: Argélia + - value: AS + label: + en_US: American Samoa + zh_Hans: 美属萨摩亚 + pt_BR: Samoa Americana + - value: AD + label: + en_US: Andorra + zh_Hans: 安道尔 + pt_BR: Andorra + - value: AO + label: + en_US: Angola + zh_Hans: 安哥拉 + pt_BR: Angola + - value: AI + label: + en_US: Anguilla + zh_Hans: 安圭拉 + pt_BR: Anguilla + - value: AQ + label: + en_US: Antarctica + zh_Hans: 南极洲 + pt_BR: Antártica + - value: AG + label: + en_US: Antigua and Barbuda + zh_Hans: 安提瓜和巴布达 + pt_BR: Antígua e Barbuda + - value: AR + label: + en_US: Argentina + zh_Hans: 阿根廷 + pt_BR: Argentina + - value: AM + label: + en_US: Armenia + zh_Hans: 亚美尼亚 + pt_BR: Armênia + - value: AW + label: + en_US: Aruba + zh_Hans: 阿鲁巴 + pt_BR: Aruba + - value: AU + label: + en_US: Australia + zh_Hans: 澳大利亚 + pt_BR: Austrália + - value: AT + label: + en_US: Austria + zh_Hans: 奥地利 + pt_BR: Áustria + - value: AZ + label: + en_US: Azerbaijan + zh_Hans: 阿塞拜疆 + pt_BR: Azerbaijão + - value: BS + label: + en_US: Bahamas + zh_Hans: 巴哈马 + pt_BR: Bahamas + - value: BH + label: + en_US: Bahrain + zh_Hans: 巴林 + pt_BR: Bahrein + - value: BD + label: + en_US: Bangladesh + zh_Hans: 孟加拉国 + pt_BR: Bangladesh + - value: BB + label: + en_US: Barbados + zh_Hans: 巴巴多斯 + pt_BR: Barbados + - value: BY + label: + en_US: Belarus + zh_Hans: 白俄罗斯 + pt_BR: Bielorrússia + - value: BE + label: + en_US: Belgium + zh_Hans: 比利时 + pt_BR: Bélgica + - value: BZ + label: + en_US: Belize + zh_Hans: 伯利兹 + pt_BR: Belize + - value: BJ + label: + en_US: Benin + zh_Hans: 贝宁 + pt_BR: Benim + - value: BM + label: + en_US: Bermuda + zh_Hans: 百慕大 + pt_BR: Bermudas + - value: BT + label: + en_US: Bhutan + zh_Hans: 不丹 + pt_BR: Butão + - value: BO + label: + en_US: Bolivia + zh_Hans: 玻利维亚 + pt_BR: Bolívia + - value: BA + label: + en_US: Bosnia and Herzegovina + zh_Hans: 波斯尼亚和黑塞哥维那 + pt_BR: Bósnia e Herzegovina + - value: BW + label: + en_US: Botswana + zh_Hans: 博茨瓦纳 + pt_BR: Botsuana + - value: BV + label: + en_US: Bouvet Island + zh_Hans: 布韦岛 + pt_BR: Ilha Bouvet + - value: BR + label: + en_US: Brazil + zh_Hans: 巴西 + pt_BR: Brasil + - value: IO + label: + en_US: British Indian Ocean Territory + zh_Hans: 英属印度洋领地 + pt_BR: Território Britânico do Oceano Índico + - value: BN + label: + en_US: Brunei Darussalam + zh_Hans: 文莱 + pt_BR: Brunei Darussalam + - value: BG + label: + en_US: Bulgaria + zh_Hans: 保加利亚 + pt_BR: Bulgária + - value: BF + label: + en_US: Burkina Faso + zh_Hans: 布基纳法索 + pt_BR: Burkina Faso + - value: BI + label: + en_US: Burundi + zh_Hans: 布隆迪 + pt_BR: Burundi + - value: KH + label: + en_US: Cambodia + zh_Hans: 柬埔寨 + pt_BR: Camboja + - value: CM + label: + en_US: Cameroon + zh_Hans: 喀麦隆 + pt_BR: Camarões + - value: CA + label: + en_US: Canada + zh_Hans: 加拿大 + pt_BR: Canadá + - value: CV + label: + en_US: Cape Verde + zh_Hans: 佛得角 + pt_BR: Cabo Verde + - value: KY + label: + en_US: Cayman Islands + zh_Hans: 开曼群岛 + pt_BR: Ilhas Cayman + - value: CF + label: + en_US: Central African Republic + zh_Hans: 中非共和国 + pt_BR: República Centro-Africana + - value: TD + label: + en_US: Chad + zh_Hans: 乍得 + pt_BR: Chade + - value: CL + label: + en_US: Chile + zh_Hans: 智利 + pt_BR: Chile + - value: CN + label: + en_US: China + zh_Hans: 中国 + pt_BR: China + - value: CX + label: + en_US: Christmas Island + zh_Hans: 圣诞岛 + pt_BR: Ilha do Natal + - value: CC + label: + en_US: Cocos (Keeling) Islands + zh_Hans: 科科斯(基林)群岛 + pt_BR: Ilhas Cocos (Keeling) + - value: CO + label: + en_US: Colombia + zh_Hans: 哥伦比亚 + pt_BR: Colômbia + - value: KM + label: + en_US: Comoros + zh_Hans: 科摩罗 + pt_BR: Comores + - value: CG + label: + en_US: Congo + zh_Hans: 刚果 + pt_BR: Congo + - value: CD + label: + en_US: Congo, the Democratic Republic of the + zh_Hans: 刚果民主共和国 + pt_BR: Congo, República Democrática do + - value: CK + label: + en_US: Cook Islands + zh_Hans: 库克群岛 + pt_BR: Ilhas Cook + - value: CR + label: + en_US: Costa Rica + zh_Hans: 哥斯达黎加 + pt_BR: Costa Rica + - value: CI + label: + en_US: Cote D'ivoire + zh_Hans: 科特迪瓦 + pt_BR: Costa do Marfim + - value: HR + label: + en_US: Croatia + zh_Hans: 克罗地亚 + pt_BR: Croácia + - value: CU + label: + en_US: Cuba + zh_Hans: 古巴 + pt_BR: Cuba + - value: CY + label: + en_US: Cyprus + zh_Hans: 塞浦路斯 + pt_BR: Chipre + - value: CZ + label: + en_US: Czech Republic + zh_Hans: 捷克共和国 + pt_BR: República Tcheca + - value: DK + label: + en_US: Denmark + zh_Hans: 丹麦 + pt_BR: Dinamarca + - value: DJ + label: + en_US: Djibouti + zh_Hans: 吉布提 + pt_BR: Djibuti + - value: DM + label: + en_US: Dominica + zh_Hans: 多米尼克 + pt_BR: Dominica + - value: DO + label: + en_US: Dominican Republic + zh_Hans: 多米尼加共和国 + pt_BR: República Dominicana + - value: EC + label: + en_US: Ecuador + zh_Hans: 厄瓜多尔 + pt_BR: Equador + - value: EG + label: + en_US: Egypt + zh_Hans: 埃及 + pt_BR: Egito + - value: SV + label: + en_US: El Salvador + zh_Hans: 萨尔瓦多 + pt_BR: El Salvador + - value: GQ + label: + en_US: Equatorial Guinea + zh_Hans: 赤道几内亚 + pt_BR: Guiné Equatorial + - value: ER + label: + en_US: Eritrea + zh_Hans: 厄立特里亚 + pt_BR: Eritreia + - value: EE + label: + en_US: Estonia + zh_Hans: 爱沙尼亚 + pt_BR: Estônia + - value: ET + label: + en_US: Ethiopia + zh_Hans: 埃塞俄比亚 + pt_BR: Etiópia + - value: FK + label: + en_US: Falkland Islands (Malvinas) + zh_Hans: 福克兰群岛(马尔维纳斯) + pt_BR: Ilhas Falkland (Malvinas) + - value: FO + label: + en_US: Faroe Islands + zh_Hans: 法罗群岛 + pt_BR: Ilhas Faroe + - value: FJ + label: + en_US: Fiji + zh_Hans: 斐济 + pt_BR: Fiji + - value: FI + label: + en_US: Finland + zh_Hans: 芬兰 + pt_BR: Finlândia + - value: FR + label: + en_US: France + zh_Hans: 法国 + pt_BR: França + - value: GF + label: + en_US: French Guiana + zh_Hans: 法属圭亚那 + pt_BR: Guiana Francesa + - value: PF + label: + en_US: French Polynesia + zh_Hans: 法属波利尼西亚 + pt_BR: Polinésia Francesa + - value: TF + label: + en_US: French Southern Territories + zh_Hans: 法属南部领地 + pt_BR: Territórios Franceses do Sul + - value: GA + label: + en_US: Gabon + zh_Hans: 加蓬 + pt_BR: Gabão + - value: GM + label: + en_US: Gambia + zh_Hans: 冈比亚 + pt_BR: Gâmbia + - value: GE + label: + en_US: Georgia + zh_Hans: 格鲁吉亚 + pt_BR: Geórgia + - value: DE + label: + en_US: Germany + zh_Hans: 德国 + pt_BR: Alemanha + - value: GH + label: + en_US: Ghana + zh_Hans: 加纳 + pt_BR: Gana + - value: GI + label: + en_US: Gibraltar + zh_Hans: 直布罗陀 + pt_BR: Gibraltar + - value: GR + label: + en_US: Greece + zh_Hans: 希腊 + pt_BR: Grécia + - value: GL + label: + en_US: Greenland + zh_Hans: 格陵兰 + pt_BR: Groenlândia + - value: GD + label: + en_US: Grenada + zh_Hans: 格林纳达 + pt_BR: Granada + - value: GP + label: + en_US: Guadeloupe + zh_Hans: 瓜德罗普 + pt_BR: Guadalupe + - value: GU + label: + en_US: Guam + zh_Hans: 关岛 + pt_BR: Guam + - value: GT + label: + en_US: Guatemala + zh_Hans: 危地马拉 + pt_BR: Guatemala + - value: GN + label: + en_US: Guinea + zh_Hans: 几内亚 + pt_BR: Guiné + - value: GW + label: + en_US: Guinea-Bissau + zh_Hans: 几内亚比绍 + pt_BR: Guiné-Bissau + - value: GY + label: + en_US: Guyana + zh_Hans: 圭亚那 + pt_BR: Guiana + - value: HT + label: + en_US: Haiti + zh_Hans: 海地 + pt_BR: Haiti + - value: HM + label: + en_US: Heard Island and McDonald Islands + zh_Hans: 赫德岛和麦克唐纳群岛 + pt_BR: Ilha Heard e Ilhas McDonald + - value: VA + label: + en_US: Holy See (Vatican City State) + zh_Hans: 教廷(梵蒂冈城国) + pt_BR: Santa Sé (Estado da Cidade do Vaticano) + - value: HN + label: + en_US: Honduras + zh_Hans: 洪都拉斯 + pt_BR: Honduras + - value: HK + label: + en_US: Hong Kong + zh_Hans: 香港 + pt_BR: Hong Kong + - value: HU + label: + en_US: Hungary + zh_Hans: 匈牙利 + pt_BR: Hungria + - value: IS + label: + en_US: Iceland + zh_Hans: 冰岛 + pt_BR: Islândia + - value: IN + label: + en_US: India + zh_Hans: 印度 + pt_BR: Índia + - value: ID + label: + en_US: Indonesia + zh_Hans: 印度尼西亚 + pt_BR: Indonésia + - value: IR + label: + en_US: Iran, Islamic Republic of + zh_Hans: 伊朗 + pt_BR: Irã + - value: IQ + label: + en_US: Iraq + zh_Hans: 伊拉克 + pt_BR: Iraque + - value: IE + label: + en_US: Ireland + zh_Hans: 爱尔兰 + pt_BR: Irlanda + - value: IL + label: + en_US: Israel + zh_Hans: 以色列 + pt_BR: Israel + - value: IT + label: + en_US: Italy + zh_Hans: 意大利 + pt_BR: Itália + - value: JM + label: + en_US: Jamaica + zh_Hans: 牙买加 + pt_BR: Jamaica + - value: JP + label: + en_US: Japan + zh_Hans: 日本 + pt_BR: Japão + - value: JO + label: + en_US: Jordan + zh_Hans: 约旦 + pt_BR: Jordânia + - value: KZ + label: + en_US: Kazakhstan + zh_Hans: 哈萨克斯坦 + pt_BR: Cazaquistão + - value: KE + label: + en_US: Kenya + zh_Hans: 肯尼亚 + pt_BR: Quênia + - value: KI + label: + en_US: Kiribati + zh_Hans: 基里巴斯 + pt_BR: Kiribati + - value: KP + label: + en_US: Korea, Democratic People's Republic of + zh_Hans: 朝鲜 + pt_BR: Coreia, República Democrática Popular da + - value: KR + label: + en_US: Korea, Republic of + zh_Hans: 韩国 + pt_BR: Coreia, República da + - value: KW + label: + en_US: Kuwait + zh_Hans: 科威特 + pt_BR: Kuwait + - value: KG + label: + en_US: Kyrgyzstan + zh_Hans: 吉尔吉斯斯坦 + pt_BR: Quirguistão + - value: LA + label: + en_US: Lao People's Democratic Republic + zh_Hans: 老挝 + pt_BR: República Democrática Popular do Laos + - value: LV + label: + en_US: Latvia + zh_Hans: 拉脱维亚 + pt_BR: Letônia + - value: LB + label: + en_US: Lebanon + zh_Hans: 黎巴嫩 + pt_BR: Líbano + - value: LS + label: + en_US: Lesotho + zh_Hans: 莱索托 + pt_BR: Lesoto + - value: LR + label: + en_US: Liberia + zh_Hans: 利比里亚 + pt_BR: Libéria + - value: LY + label: + en_US: Libyan Arab Jamahiriya + zh_Hans: 利比亚 + pt_BR: Líbia + - value: LI + label: + en_US: Liechtenstein + zh_Hans: 列支敦士登 + pt_BR: Liechtenstein + - value: LT + label: + en_US: Lithuania + zh_Hans: 立陶宛 + pt_BR: Lituânia + - value: LU + label: + en_US: Luxembourg + zh_Hans: 卢森堡 + pt_BR: Luxemburgo + - value: MO + label: + en_US: Macao + zh_Hans: 澳门 + pt_BR: Macau + - value: MK + label: + en_US: Macedonia, the Former Yugosalv Republic of + zh_Hans: 前南斯拉夫马其顿共和国 + pt_BR: Macedônia, Ex-República Iugoslava da + - value: MG + label: + en_US: Madagascar + zh_Hans: 马达加斯加 + pt_BR: Madagascar + - value: MW + label: + en_US: Malawi + zh_Hans: 马拉维 + pt_BR: Malaui + - value: MY + label: + en_US: Malaysia + zh_Hans: 马来西亚 + pt_BR: Malásia + - value: MV + label: + en_US: Maldives + zh_Hans: 马尔代夫 + pt_BR: Maldivas + - value: ML + label: + en_US: Mali + zh_Hans: 马里 + pt_BR: Mali + - value: MT + label: + en_US: Malta + zh_Hans: 马耳他 + pt_BR: Malta + - value: MH + label: + en_US: Marshall Islands + zh_Hans: 马绍尔群岛 + pt_BR: Ilhas Marshall + - value: MQ + label: + en_US: Martinique + zh_Hans: 马提尼克 + pt_BR: Martinica + - value: MR + label: + en_US: Mauritania + zh_Hans: 毛里塔尼亚 + pt_BR: Mauritânia + - value: MU + label: + en_US: Mauritius + zh_Hans: 毛里求斯 + pt_BR: Maurício + - value: YT + label: + en_US: Mayotte + zh_Hans: 马约特 + pt_BR: Mayotte + - value: MX + label: + en_US: Mexico + zh_Hans: 墨西哥 + pt_BR: México + - value: FM + label: + en_US: Micronesia, Federated States of + zh_Hans: 密克罗尼西亚联邦 + pt_BR: Micronésia, Estados Federados da + - value: MD + label: + en_US: Moldova, Republic of + zh_Hans: 摩尔多瓦共和国 + pt_BR: Moldávia, República da + - value: MC + label: + en_US: Monaco + zh_Hans: 摩纳哥 + pt_BR: Mônaco + - value: MN + label: + en_US: Mongolia + zh_Hans: 蒙古 + pt_BR: Mongólia + - value: MS + label: + en_US: Montserrat + zh_Hans: 蒙特塞拉特 + pt_BR: Montserrat + - value: MA + label: + en_US: Morocco + zh_Hans: 摩洛哥 + pt_BR: Marrocos + - value: MZ + label: + en_US: Mozambique + zh_Hans: 莫桑比克 + pt_BR: Moçambique + - value: MM + label: + en_US: Myanmar + zh_Hans: 缅甸 + pt_BR: Mianmar + - value: NA + label: + en_US: Namibia + zh_Hans: 纳米比亚 + pt_BR: Namíbia + - value: NR + label: + en_US: Nauru + zh_Hans: 瑙鲁 + pt_BR: Nauru + - value: NP + label: + en_US: Nepal + zh_Hans: 尼泊尔 + pt_BR: Nepal + - value: NL + label: + en_US: Netherlands + zh_Hans: 荷兰 + pt_BR: Países Baixos + - value: AN + label: + en_US: Netherlands Antilles + zh_Hans: 荷属安的列斯 + pt_BR: Antilhas Holandesas + - value: NC + label: + en_US: New Caledonia + zh_Hans: 新喀里多尼亚 + pt_BR: Nova Caledônia + - value: NZ + label: + en_US: New Zealand + zh_Hans: 新西兰 + pt_BR: Nova Zelândia + - value: NI + label: + en_US: Nicaragua + zh_Hans: 尼加拉瓜 + pt_BR: Nicarágua + - value: NE + label: + en_US: Niger + zh_Hans: 尼日尔 + pt_BR: Níger + - value: NG + label: + en_US: Nigeria + zh_Hans: 尼日利亚 + pt_BR: Nigéria + - value: NU + label: + en_US: Niue + zh_Hans: 纽埃 + pt_BR: Niue + - value: NF + label: + en_US: Norfolk Island + zh_Hans: 诺福克岛 + pt_BR: Ilha Norfolk + - value: MP + label: + en_US: Northern Mariana Islands + zh_Hans: 北马里亚纳群岛 + pt_BR: Ilhas Marianas do Norte + - value: "NO" + label: + en_US: Norway + zh_Hans: 挪威 + pt_BR: Noruega + - value: OM + label: + en_US: Oman + zh_Hans: 阿曼 + pt_BR: Omã + - value: PK + label: + en_US: Pakistan + zh_Hans: 巴基斯坦 + pt_BR: Paquistão + - value: PW + label: + en_US: Palau + zh_Hans: 帕劳 + pt_BR: Palau + - value: PS + label: + en_US: Palestinian Territory, Occupied + zh_Hans: 巴勒斯坦领土 + pt_BR: Palestina, Território Ocupado + - value: PA + label: + en_US: Panama + zh_Hans: 巴拿马 + pt_BR: Panamá + - value: PG + label: + en_US: Papua New Guinea + zh_Hans: 巴布亚新几内亚 + pt_BR: Papua Nova Guiné + - value: PY + label: + en_US: Paraguay + zh_Hans: 巴拉圭 + pt_BR: Paraguai + - value: PE + label: + en_US: Peru + zh_Hans: 秘鲁 + pt_BR: Peru + - value: PH + label: + en_US: Philippines + zh_Hans: 菲律宾 + pt_BR: Filipinas + - value: PN + label: + en_US: Pitcairn + zh_Hans: 皮特凯恩岛 + pt_BR: Pitcairn + - value: PL + label: + en_US: Poland + zh_Hans: 波兰 + pt_BR: Polônia + - value: PT + label: + en_US: Portugal + zh_Hans: 葡萄牙 + pt_BR: Portugal + - value: PR + label: + en_US: Puerto Rico + zh_Hans: 波多黎各 + pt_BR: Porto Rico + - value: QA + label: + en_US: Qatar + zh_Hans: 卡塔尔 + pt_BR: Catar + - value: RE + label: + en_US: Reunion + zh_Hans: 留尼旺 + pt_BR: Reunião + - value: RO + label: + en_US: Romania + zh_Hans: 罗马尼亚 + pt_BR: Romênia + - value: RU + label: + en_US: Russian Federation + zh_Hans: 俄罗斯联邦 + pt_BR: Rússia + - value: RW + label: + en_US: Rwanda + zh_Hans: 卢旺达 + pt_BR: Ruanda + - value: SH + label: + en_US: Saint Helena + zh_Hans: 圣赫勒拿 + pt_BR: Santa Helena + - value: KN + label: + en_US: Saint Kitts and Nevis + zh_Hans: 圣基茨和尼维斯 + pt_BR: São Cristóvão e Nevis + - value: LC + label: + en_US: Saint Lucia + zh_Hans: 圣卢西亚 + pt_BR: Santa Lúcia + - value: PM + label: + en_US: Saint Pierre and Miquelon + zh_Hans: 圣皮埃尔和密克隆 + pt_BR: São Pedro e Miquelon + - value: VC + label: + en_US: Saint Vincent and the Grenadines + zh_Hans: 圣文森特和格林纳丁斯 + pt_BR: São Vicente e Granadinas + - value: WS + label: + en_US: Samoa + zh_Hans: 萨摩亚 + pt_BR: Samoa + - value: SM + label: + en_US: San Marino + zh_Hans: 圣马力诺 + pt_BR: San Marino + - value: ST + label: + en_US: Sao Tome and Principe + zh_Hans: 圣多美和普林西比 + pt_BR: São Tomé e Príncipe + - value: SA + label: + en_US: Saudi Arabia + zh_Hans: 沙特阿拉伯 + pt_BR: Arábia Saudita + - value: SN + label: + en_US: Senegal + zh_Hans: 塞内加尔 + pt_BR: Senegal + - value: RS + label: + en_US: Serbia and Montenegro + zh_Hans: 塞尔维亚和黑山 + pt_BR: Sérvia e Montenegro + - value: SC + label: + en_US: Seychelles + zh_Hans: 塞舌尔 + pt_BR: Seicheles + - value: SL + label: + en_US: Sierra Leone + zh_Hans: 塞拉利昂 + pt_BR: Serra Leoa + - value: SG + label: + en_US: Singapore + zh_Hans: 新加坡 + pt_BR: Singapura + - value: SK + label: + en_US: Slovakia + zh_Hans: 斯洛伐克 + pt_BR: Eslováquia + - value: SI + label: + en_US: Slovenia + zh_Hans: 斯洛文尼亚 + pt_BR: Eslovênia + - value: SB + label: + en_US: Solomon Islands + zh_Hans: 所罗门群岛 + pt_BR: Ilhas Salomão + - value: SO + label: + en_US: Somalia + zh_Hans: 索马里 + pt_BR: Somália + - value: ZA + label: + en_US: South Africa + zh_Hans: 南非 + pt_BR: África do Sul + - value: GS + label: + en_US: South Georgia and the South Sandwich Islands + zh_Hans: 南乔治亚和南桑威奇群岛 + pt_BR: Geórgia do Sul e Ilhas Sandwich do Sul + - value: ES + label: + en_US: Spain + zh_Hans: 西班牙 + pt_BR: Espanha + - value: LK + label: + en_US: Sri Lanka + zh_Hans: 斯里兰卡 + pt_BR: Sri Lanka + - value: SD + label: + en_US: Sudan + zh_Hans: 苏丹 + pt_BR: Sudão + - value: SR + label: + en_US: Suriname + zh_Hans: 苏里南 + pt_BR: Suriname + - value: SJ + label: + en_US: Svalbard and Jan Mayen + zh_Hans: 斯瓦尔巴特和扬马延岛 + pt_BR: Svalbard e Jan Mayen + - value: SZ + label: + en_US: Swaziland + zh_Hans: 斯威士兰 + pt_BR: Essuatíni + - value: SE + label: + en_US: Sweden + zh_Hans: 瑞典 + pt_BR: Suécia + - value: CH + label: + en_US: Switzerland + zh_Hans: 瑞士 + pt_BR: Suíça + - value: SY + label: + en_US: Syrian Arab Republic + zh_Hans: 叙利亚 + pt_BR: Síria + - value: TW + label: + en_US: Taiwan, Province of China + zh_Hans: 台湾 + pt_BR: Taiwan + - value: TJ + label: + en_US: Tajikistan + zh_Hans: 塔吉克斯坦 + pt_BR: Tajiquistão + - value: TZ + label: + en_US: Tanzania, United Republic of + zh_Hans: 坦桑尼亚联合共和国 + pt_BR: Tanzânia + - value: TH + label: + en_US: Thailand + zh_Hans: 泰国 + pt_BR: Tailândia + - value: TL + label: + en_US: Timor-Leste + zh_Hans: 东帝汶 + pt_BR: Timor-Leste + - value: TG + label: + en_US: Togo + zh_Hans: 多哥 + pt_BR: Togo + - value: TK + label: + en_US: Tokelau + zh_Hans: 托克劳 + pt_BR: Toquelau + - value: TO + label: + en_US: Tonga + zh_Hans: 汤加 + pt_BR: Tonga + - value: TT + label: + en_US: Trinidad and Tobago + zh_Hans: 特立尼达和多巴哥 + pt_BR: Trindade e Tobago + - value: TN + label: + en_US: Tunisia + zh_Hans: 突尼斯 + pt_BR: Tunísia + - value: TR + label: + en_US: Turkey + zh_Hans: 土耳其 + pt_BR: Turquia + - value: TM + label: + en_US: Turkmenistan + zh_Hans: 土库曼斯坦 + pt_BR: Turcomenistão + - value: TC + label: + en_US: Turks and Caicos Islands + zh_Hans: 特克斯和凯科斯群岛 + pt_BR: Ilhas Turks e Caicos + - value: TV + label: + en_US: Tuvalu + zh_Hans: 图瓦卢 + pt_BR: Tuvalu + - value: UG + label: + en_US: Uganda + zh_Hans: 乌干达 + pt_BR: Uganda + - value: UA + label: + en_US: Ukraine + zh_Hans: 乌克兰 + pt_BR: Ucrânia + - value: AE + label: + en_US: United Arab Emirates + zh_Hans: 阿联酋 + pt_BR: Emirados Árabes Unidos + - value: UK + label: + en_US: United Kingdom + zh_Hans: 英国 + pt_BR: Reino Unido + - value: GB + label: + en_US: United Kingdom + zh_Hans: 英国 + pt_BR: Reino Unido + - value: US + label: + en_US: United States + zh_Hans: 美国 + pt_BR: Estados Unidos + - value: UM + label: + en_US: United States Minor Outlying Islands + zh_Hans: 美国本土外小岛屿 + pt_BR: Ilhas Menores Distantes dos Estados Unidos + - value: UY + label: + en_US: Uruguay + zh_Hans: 乌拉圭 + pt_BR: Uruguai + - value: UZ + label: + en_US: Uzbekistan + zh_Hans: 乌兹别克斯坦 + pt_BR: Uzbequistão + - value: VU + label: + en_US: Vanuatu + zh_Hans: 瓦努阿图 + pt_BR: Vanuatu + - value: VE + label: + en_US: Venezuela + zh_Hans: 委内瑞拉 + pt_BR: Venezuela + - value: VN + label: + en_US: Viet Nam + zh_Hans: 越南 + pt_BR: Vietnã + - value: VG + label: + en_US: Virgin Islands, British + zh_Hans: 英属维尔京群岛 + pt_BR: Ilhas Virgens Britânicas + - value: VI + label: + en_US: Virgin Islands, U.S. + zh_Hans: 美属维尔京群岛 + pt_BR: Ilhas Virgens dos EUA + - value: WF + label: + en_US: Wallis and Futuna + zh_Hans: 瓦利斯和富图纳群岛 + pt_BR: Wallis e Futuna + - value: EH + label: + en_US: Western Sahara + zh_Hans: 西撒哈拉 + pt_BR: Saara Ocidental + - value: YE + label: + en_US: Yemen + zh_Hans: 也门 + pt_BR: Iémen + - value: ZM + label: + en_US: Zambia + zh_Hans: 赞比亚 + pt_BR: Zâmbia + - value: ZW + label: + en_US: Zimbabwe + zh_Hans: 津巴布韦 + pt_BR: Zimbábue + - name: hl + type: select + label: + en_US: Language + zh_Hans: 语言 + human_description: + en_US: Defines the interface language of the search. Default is "en". + zh_Hans: 定义搜索的界面语言。默认为“en”。 + required: false + default: en + form: form + options: + - value: af + label: + en_US: Afrikaans + zh_Hans: 南非语 + - value: ak + label: + en_US: Akan + zh_Hans: 阿坎语 + - value: sq + label: + en_US: Albanian + zh_Hans: 阿尔巴尼亚语 + - value: ws + label: + en_US: Samoa + zh_Hans: 萨摩亚语 + - value: am + label: + en_US: Amharic + zh_Hans: 阿姆哈拉语 + - value: ar + label: + en_US: Arabic + zh_Hans: 阿拉伯语 + - value: hy + label: + en_US: Armenian + zh_Hans: 亚美尼亚语 + - value: az + label: + en_US: Azerbaijani + zh_Hans: 阿塞拜疆语 + - value: eu + label: + en_US: Basque + zh_Hans: 巴斯克语 + - value: be + label: + en_US: Belarusian + zh_Hans: 白俄罗斯语 + - value: bem + label: + en_US: Bemba + zh_Hans: 班巴语 + - value: bn + label: + en_US: Bengali + zh_Hans: 孟加拉语 + - value: bh + label: + en_US: Bihari + zh_Hans: 比哈尔语 + - value: xx-bork + label: + en_US: Bork, bork, bork! + zh_Hans: 博克语 + - value: bs + label: + en_US: Bosnian + zh_Hans: 波斯尼亚语 + - value: br + label: + en_US: Breton + zh_Hans: 布列塔尼语 + - value: bg + label: + en_US: Bulgarian + zh_Hans: 保加利亚语 + - value: bt + label: + en_US: Bhutanese + zh_Hans: 不丹语 + - value: km + label: + en_US: Cambodian + zh_Hans: 高棉语 + - value: ca + label: + en_US: Catalan + zh_Hans: 加泰罗尼亚语 + - value: chr + label: + en_US: Cherokee + zh_Hans: 切罗基语 + - value: ny + label: + en_US: Chichewa + zh_Hans: 齐切瓦语 + - value: zh-cn + label: + en_US: Chinese (Simplified) + zh_Hans: 中文(简体) + - value: zh-tw + label: + en_US: Chinese (Traditional) + zh_Hans: 中文(繁体) + - value: co + label: + en_US: Corsican + zh_Hans: 科西嘉语 + - value: hr + label: + en_US: Croatian + zh_Hans: 克罗地亚语 + - value: cs + label: + en_US: Czech + zh_Hans: 捷克语 + - value: da + label: + en_US: Danish + zh_Hans: 丹麦语 + - value: nl + label: + en_US: Dutch + zh_Hans: 荷兰语 + - value: xx-elmer + label: + en_US: Elmer Fudd + zh_Hans: 艾尔默福德语 + - value: en + label: + en_US: English + zh_Hans: 英语 + - value: eo + label: + en_US: Esperanto + zh_Hans: 世界语 + - value: et + label: + en_US: Estonian + zh_Hans: 爱沙尼亚语 + - value: ee + label: + en_US: Ewe + zh_Hans: 埃维语 + - value: fo + label: + en_US: Faroese + zh_Hans: 法罗语 + - value: tl + label: + en_US: Filipino + zh_Hans: 菲律宾语 + - value: fi + label: + en_US: Finnish + zh_Hans: 芬兰语 + - value: fr + label: + en_US: French + zh_Hans: 法语 + - value: fy + label: + en_US: Frisian + zh_Hans: 弗里西亚语 + - value: gaa + label: + en_US: Ga + zh_Hans: 加语 + - value: gl + label: + en_US: Galician + zh_Hans: 加利西亚语 + - value: ka + label: + en_US: Georgian + zh_Hans: 格鲁吉亚语 + - value: de + label: + en_US: German + zh_Hans: 德语 + - value: el + label: + en_US: Greek + zh_Hans: 希腊语 + - value: kl + label: + en_US: Greenlandic + zh_Hans: 格陵兰语 + - value: gn + label: + en_US: Guarani + zh_Hans: 瓜拉尼语 + - value: gu + label: + en_US: Gujarati + zh_Hans: 古吉拉特语 + - value: xx-hacker + label: + en_US: Hacker + zh_Hans: 黑客语 + - value: ht + label: + en_US: Haitian Creole + zh_Hans: 海地克里奥尔语 + - value: ha + label: + en_US: Hausa + zh_Hans: 豪萨语 + - value: haw + label: + en_US: Hawaiian + zh_Hans: 夏威夷语 + - value: iw + label: + en_US: Hebrew + zh_Hans: 希伯来语 + - value: hi + label: + en_US: Hindi + zh_Hans: 印地语 + - value: hu + label: + en_US: Hungarian + zh_Hans: 匈牙利语 + - value: is + label: + en_US: Icelandic + zh_Hans: 冰岛语 + - value: ig + label: + en_US: Igbo + zh_Hans: 伊博语 + - value: id + label: + en_US: Indonesian + zh_Hans: 印尼语 + - value: ia + label: + en_US: Interlingua + zh_Hans: 国际语 + - value: ga + label: + en_US: Irish + zh_Hans: 爱尔兰语 + - value: it + label: + en_US: Italian + zh_Hans: 意大利语 + - value: ja + label: + en_US: Japanese + zh_Hans: 日语 + - value: jw + label: + en_US: Javanese + zh_Hans: 爪哇语 + - value: kn + label: + en_US: Kannada + zh_Hans: 卡纳达语 + - value: kk + label: + en_US: Kazakh + zh_Hans: 哈萨克语 + - value: rw + label: + en_US: Kinyarwanda + zh_Hans: 基尼亚卢旺达语 + - value: rn + label: + en_US: Kirundi + zh_Hans: 基隆迪语 + - value: xx-klingon + label: + en_US: Klingon + zh_Hans: 克林贡语 + - value: kg + label: + en_US: Kongo + zh_Hans: 刚果语 + - value: ko + label: + en_US: Korean + zh_Hans: 韩语 + - value: kri + label: + en_US: Krio (Sierra Leone) + zh_Hans: 塞拉利昂克里奥尔语 + - value: ku + label: + en_US: Kurdish + zh_Hans: 库尔德语 + - value: ckb + label: + en_US: Kurdish (Soranî) + zh_Hans: 库尔德语(索拉尼) + - value: ky + label: + en_US: Kyrgyz + zh_Hans: 吉尔吉斯语 + - value: lo + label: + en_US: Laothian + zh_Hans: 老挝语 + - value: la + label: + en_US: Latin + zh_Hans: 拉丁语 + - value: lv + label: + en_US: Latvian + zh_Hans: 拉脱维亚语 + - value: ln + label: + en_US: Lingala + zh_Hans: 林加拉语 + - value: lt + label: + en_US: Lithuanian + zh_Hans: 立陶宛语 + - value: loz + label: + en_US: Lozi + zh_Hans: 洛齐语 + - value: lg + label: + en_US: Luganda + zh_Hans: 卢干达语 + - value: ach + label: + en_US: Luo + zh_Hans: 卢奥语 + - value: mk + label: + en_US: Macedonian + zh_Hans: 马其顿语 + - value: mg + label: + en_US: Malagasy + zh_Hans: 马尔加什语 + - value: my + label: + en_US: Malay + zh_Hans: 马来语 + - value: ml + label: + en_US: Malayalam + zh_Hans: 马拉雅拉姆语 + - value: mt + label: + en_US: Maltese + zh_Hans: 马耳他语 + - value: mv + label: + en_US: Maldives + zh_Hans: 马尔代夫语 + - value: mi + label: + en_US: Maori + zh_Hans: 毛利语 + - value: mr + label: + en_US: Marathi + zh_Hans: 马拉地语 + - value: mfe + label: + en_US: Mauritian Creole + zh_Hans: 毛里求斯克里奥尔语 + - value: mo + label: + en_US: Moldavian + zh_Hans: 摩尔达维亚语 + - value: mn + label: + en_US: Mongolian + zh_Hans: 蒙古语 + - value: sr-me + label: + en_US: Montenegrin + zh_Hans: 黑山语 + - value: ne + label: + en_US: Nepali + zh_Hans: 尼泊尔语 + - value: pcm + label: + en_US: Nigerian Pidgin + zh_Hans: 尼日利亚皮钦语 + - value: nso + label: + en_US: Northern Sotho + zh_Hans: 北索托语 + - value: "no" + label: + en_US: Norwegian + zh_Hans: 挪威语 + - value: nn + label: + en_US: Norwegian (Nynorsk) + zh_Hans: 挪威语(尼诺斯克语) + - value: oc + label: + en_US: Occitan + zh_Hans: 奥克语 + - value: or + label: + en_US: Oriya + zh_Hans: 奥里亚语 + - value: om + label: + en_US: Oromo + zh_Hans: 奥罗莫语 + - value: ps + label: + en_US: Pashto + zh_Hans: 普什图语 + - value: fa + label: + en_US: Persian + zh_Hans: 波斯语 + - value: xx-pirate + label: + en_US: Pirate + zh_Hans: 海盗语 + - value: pl + label: + en_US: Polish + zh_Hans: 波兰语 + - value: pt + label: + en_US: Portuguese + zh_Hans: 葡萄牙语 + - value: pt-br + label: + en_US: Portuguese (Brazil) + zh_Hans: 葡萄牙语(巴西) + - value: pt-pt + label: + en_US: Portuguese (Portugal) + zh_Hans: 葡萄牙语(葡萄牙) + - value: pa + label: + en_US: Punjabi + zh_Hans: 旁遮普语 + - value: qu + label: + en_US: Quechua + zh_Hans: 克丘亚语 + - value: ro + label: + en_US: Romanian + zh_Hans: 罗马尼亚语 + - value: rm + label: + en_US: Romansh + zh_Hans: 罗曼什语 + - value: nyn + label: + en_US: Runyakitara + zh_Hans: 卢尼亚基塔拉语 + - value: ru + label: + en_US: Russian + zh_Hans: 俄语 + - value: gd + label: + en_US: Scots Gaelic + zh_Hans: 苏格兰盖尔语 + - value: sr + label: + en_US: Serbian + zh_Hans: 塞尔维亚语 + - value: sh + label: + en_US: Serbo-Croatian + zh_Hans: 塞尔维亚-克罗地亚语 + - value: st + label: + en_US: Sesotho + zh_Hans: 塞索托语 + - value: tn + label: + en_US: Setswana + zh_Hans: 塞茨瓦纳语 + - value: crs + label: + en_US: Seychellois Creole + zh_Hans: 塞舌尔克里奥尔语 + - value: sn + label: + en_US: Shona + zh_Hans: 绍纳语 + - value: sd + label: + en_US: Sindhi + zh_Hans: 信德语 + - value: si + label: + en_US: Sinhalese + zh_Hans: 僧伽罗语 + - value: sk + label: + en_US: Slovak + zh_Hans: 斯洛伐克语 + - value: sl + label: + en_US: Slovenian + zh_Hans: 斯洛文尼亚语 + - value: so + label: + en_US: Somali + zh_Hans: 索马里语 + - value: es + label: + en_US: Spanish + zh_Hans: 西班牙语 + - value: es-419 + label: + en_US: Spanish (Latin American) + zh_Hans: 西班牙语(拉丁美洲) + - value: su + label: + en_US: Sundanese + zh_Hans: 巽他语 + - value: sw + label: + en_US: Swahili + zh_Hans: 斯瓦希里语 + - value: sv + label: + en_US: Swedish + zh_Hans: 瑞典语 + - value: tg + label: + en_US: Tajik + zh_Hans: 塔吉克语 + - value: ta + label: + en_US: Tamil + zh_Hans: 泰米尔语 + - value: tt + label: + en_US: Tatar + zh_Hans: 鞑靼语 + - value: te + label: + en_US: Telugu + zh_Hans: 泰卢固语 + - value: th + label: + en_US: Thai + zh_Hans: 泰语 + - value: ti + label: + en_US: Tigrinya + zh_Hans: 提格利尼亚语 + - value: to + label: + en_US: Tonga + zh_Hans: 汤加语 + - value: lua + label: + en_US: Tshiluba + zh_Hans: 卢巴语 + - value: tum + label: + en_US: Tumbuka + zh_Hans: 图布卡语 + - value: tr + label: + en_US: Turkish + zh_Hans: 土耳其语 + - value: tk + label: + en_US: Turkmen + zh_Hans: 土库曼语 + - value: tw + label: + en_US: Twi + zh_Hans: 契维语 + - value: ug + label: + en_US: Uighur + zh_Hans: 维吾尔语 + - value: uk + label: + en_US: Ukrainian + zh_Hans: 乌克兰语 + - value: ur + label: + en_US: Urdu + zh_Hans: 乌尔都语 + - value: uz + label: + en_US: Uzbek + zh_Hans: 乌兹别克语 + - value: vu + label: + en_US: Vanuatu + zh_Hans: 瓦努阿图语 + - value: vi + label: + en_US: Vietnamese + zh_Hans: 越南语 + - value: cy + label: + en_US: Welsh + zh_Hans: 威尔士语 + - value: wo + label: + en_US: Wolof + zh_Hans: 沃洛夫语 + - value: xh + label: + en_US: Xhosa + zh_Hans: 科萨语 + - value: yi + label: + en_US: Yiddish + zh_Hans: 意第绪语 + - value: yo + label: + en_US: Yoruba + zh_Hans: 约鲁巴语 + - value: zu + label: + en_US: Zulu + zh_Hans: 祖鲁语 + - name: google_domain + type: string + required: false + label: + en_US: google_domain + zh_Hans: google_domain + human_description: + en_US: Defines the Google domain of the search. Default is "google.com". + zh_Hans: 定义搜索的 Google 域。默认为“google.com”。 + llm_description: Defines Google domain in which you want to search. + form: llm + - name: num + type: number + required: false + label: + en_US: num + zh_Hans: num + human_description: + en_US: Specifies the number of results to display per page. Default is 10. Max number - 100, min - 1. + zh_Hans: 指定每页显示的结果数。默认值为 10。最大数量 - 100,最小数量 - 1。 + llm_description: Specifies the num of results to display per page. + form: llm diff --git a/api/core/tools/provider/builtin/searchapi/tools/google_jobs.py b/api/core/tools/provider/builtin/searchapi/tools/google_jobs.py new file mode 100644 index 0000000000000000000000000000000000000000..de42360898b7e0893158d39411f135ccc42a54b1 --- /dev/null +++ b/api/core/tools/provider/builtin/searchapi/tools/google_jobs.py @@ -0,0 +1,102 @@ +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +SEARCH_API_URL = "https://www.searchapi.io/api/v1/search" + + +class SearchAPI: + """ + SearchAPI tool provider. + """ + + def __init__(self, api_key: str) -> None: + """Initialize SearchAPI tool provider.""" + self.searchapi_api_key = api_key + + def run(self, query: str, **kwargs: Any) -> str: + """Run query through SearchAPI and parse result.""" + type = kwargs.get("result_type", "text") + return self._process_response(self.results(query, **kwargs), type=type) + + def results(self, query: str, **kwargs: Any) -> dict: + """Run query through SearchAPI and return the raw result.""" + params = self.get_params(query, **kwargs) + response = requests.get( + url=SEARCH_API_URL, + params=params, + headers={"Authorization": f"Bearer {self.searchapi_api_key}"}, + ) + response.raise_for_status() + return response.json() + + def get_params(self, query: str, **kwargs: Any) -> dict[str, str]: + """Get parameters for SearchAPI.""" + return { + "engine": "google_jobs", + "q": query, + **{key: value for key, value in kwargs.items() if value not in {None, ""}}, + } + + @staticmethod + def _process_response(res: dict, type: str) -> str: + """Process response from SearchAPI.""" + if "error" in res: + return res["error"] + + toret = "" + if type == "text": + if "jobs" in res and "title" in res["jobs"][0]: + for item in res["jobs"]: + toret += ( + "title: " + + item["title"] + + "\n" + + "company_name: " + + item["company_name"] + + "content: " + + item["description"] + + "\n" + ) + if toret == "": + toret = "No good search result found" + + elif type == "link": + if "jobs" in res and "apply_link" in res["jobs"][0]: + for item in res["jobs"]: + toret += f"[{item['title']} - {item['company_name']}]({item['apply_link']})\n" + else: + toret = "No good search result found" + return toret + + +class GoogleJobsTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Invoke the SearchApi tool. + """ + query = tool_parameters["query"] + result_type = tool_parameters["result_type"] + is_remote = tool_parameters.get("is_remote") + google_domain = tool_parameters.get("google_domain", "google.com") + gl = tool_parameters.get("gl", "us") + hl = tool_parameters.get("hl", "en") + location = tool_parameters.get("location") + + ltype = 1 if is_remote else None + + api_key = self.runtime.credentials["searchapi_api_key"] + result = SearchAPI(api_key).run( + query, result_type=result_type, google_domain=google_domain, gl=gl, hl=hl, location=location, ltype=ltype + ) + + if result_type == "text": + return self.create_text_message(text=result) + return self.create_link_message(link=result) diff --git a/api/core/tools/provider/builtin/searchapi/tools/google_jobs.yaml b/api/core/tools/provider/builtin/searchapi/tools/google_jobs.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3e00e20fbd6e33caef0171dbb8d03e4bf29fa2db --- /dev/null +++ b/api/core/tools/provider/builtin/searchapi/tools/google_jobs.yaml @@ -0,0 +1,1403 @@ +identity: + name: google_jobs_api + author: SearchApi + label: + en_US: Google Jobs API + zh_Hans: Google Jobs API +description: + human: + en_US: A tool to retrieve job titles, company names and description from Google Jobs engine. + zh_Hans: 一个从 Google 招聘引擎检索职位名称、公司名称和描述的工具。 + llm: A tool to retrieve job titles, company names and description from Google Jobs engine. +parameters: + - name: query + type: string + required: true + label: + en_US: Query + zh_Hans: 询问 + human_description: + en_US: Defines the query you want to search. + zh_Hans: 定义您要搜索的查询。 + llm_description: Defines the search query you want to search. + form: llm + - name: result_type + type: select + required: true + options: + - value: text + label: + en_US: text + zh_Hans: 文本 + - value: link + label: + en_US: link + zh_Hans: 链接 + default: text + label: + en_US: Result type + zh_Hans: 结果类型 + human_description: + en_US: used for selecting the result type, text or link + zh_Hans: 用于选择结果类型,使用文本还是链接进行展示 + form: form + - name: location + type: string + required: false + label: + en_US: Location + zh_Hans: 询问 + human_description: + en_US: Defines from where you want the search to originate. (For example - New York) + zh_Hans: 定义您想要搜索的起始位置。 (例如 - 纽约) + llm_description: Defines from where you want the search to originate. (For example - New York) + form: llm + - name: gl + type: select + label: + en_US: Country + zh_Hans: 国家 + required: false + human_description: + en_US: Defines the country of the search. Default is "US". + zh_Hans: 定义搜索的国家/地区。默认为“美国”。 + llm_description: Defines the gl parameter of the Google search. + form: form + default: US + options: + - value: DZ + label: + en_US: Algeria + zh_Hans: 阿尔及利亚 + pt_BR: Algeria + - value: AS + label: + en_US: American Samoa + zh_Hans: 美属萨摩亚 + pt_BR: American Samoa + - value: AO + label: + en_US: Angola + zh_Hans: 安哥拉 + pt_BR: Angola + - value: AI + label: + en_US: Anguilla + zh_Hans: 安圭拉 + pt_BR: Anguilla + - value: AG + label: + en_US: Antigua and Barbuda + zh_Hans: 安提瓜和巴布达 + pt_BR: Antigua and Barbuda + - value: AW + label: + en_US: Aruba + zh_Hans: 阿鲁巴 + pt_BR: Aruba + - value: AT + label: + en_US: Austria + zh_Hans: 奥地利 + pt_BR: Austria + - value: BS + label: + en_US: Bahamas + zh_Hans: 巴哈马 + pt_BR: Bahamas + - value: BH + label: + en_US: Bahrain + zh_Hans: 巴林 + pt_BR: Bahrain + - value: BD + label: + en_US: Bangladesh + zh_Hans: 孟加拉国 + pt_BR: Bangladesh + - value: BY + label: + en_US: Belarus + zh_Hans: 白俄罗斯 + pt_BR: Belarus + - value: BE + label: + en_US: Belgium + zh_Hans: 比利时 + pt_BR: Belgium + - value: BZ + label: + en_US: Belize + zh_Hans: 伯利兹 + pt_BR: Belize + - value: BJ + label: + en_US: Benin + zh_Hans: 贝宁 + pt_BR: Benin + - value: BM + label: + en_US: Bermuda + zh_Hans: 百慕大 + pt_BR: Bermuda + - value: BO + label: + en_US: Bolivia + zh_Hans: 玻利维亚 + pt_BR: Bolivia + - value: BW + label: + en_US: Botswana + zh_Hans: 博茨瓦纳 + pt_BR: Botswana + - value: BR + label: + en_US: Brazil + zh_Hans: 巴西 + pt_BR: Brazil + - value: IO + label: + en_US: British Indian Ocean Territory + zh_Hans: 英属印度洋领地 + pt_BR: British Indian Ocean Territory + - value: BF + label: + en_US: Burkina Faso + zh_Hans: 布基纳法索 + pt_BR: Burkina Faso + - value: BI + label: + en_US: Burundi + zh_Hans: 布隆迪 + pt_BR: Burundi + - value: CM + label: + en_US: Cameroon + zh_Hans: 喀麦隆 + pt_BR: Cameroon + - value: CA + label: + en_US: Canada + zh_Hans: 加拿大 + pt_BR: Canada + - value: CV + label: + en_US: Cape Verde + zh_Hans: 佛得角 + pt_BR: Cape Verde + - value: KY + label: + en_US: Cayman Islands + zh_Hans: 开曼群岛 + pt_BR: Cayman Islands + - value: CF + label: + en_US: Central African Republic + zh_Hans: 中非共和国 + pt_BR: Central African Republic + - value: TD + label: + en_US: Chad + zh_Hans: 乍得 + pt_BR: Chad + - value: CL + label: + en_US: Chile + zh_Hans: 智利 + pt_BR: Chile + - value: CO + label: + en_US: Colombia + zh_Hans: 哥伦比亚 + pt_BR: Colombia + - value: CD + label: + en_US: Congo, the Democratic Republic of the + zh_Hans: 刚果民主共和国 + pt_BR: Congo, the Democratic Republic of the + - value: CR + label: + en_US: Costa Rica + zh_Hans: 哥斯达黎加 + pt_BR: Costa Rica + - value: CI + label: + en_US: Cote D'ivoire + zh_Hans: 科特迪瓦 + pt_BR: Cote D'ivoire + - value: CU + label: + en_US: Cuba + zh_Hans: 古巴 + pt_BR: Cuba + - value: DK + label: + en_US: Denmark + zh_Hans: 丹麦 + pt_BR: Denmark + - value: DJ + label: + en_US: Djibouti + zh_Hans: 吉布提 + pt_BR: Djibouti + - value: DM + label: + en_US: Dominica + zh_Hans: 多米尼克 + pt_BR: Dominica + - value: DO + label: + en_US: Dominican Republic + zh_Hans: 多米尼加共和国 + pt_BR: Dominican Republic + - value: EC + label: + en_US: Ecuador + zh_Hans: 厄瓜多尔 + pt_BR: Ecuador + - value: EG + label: + en_US: Egypt + zh_Hans: 埃及 + pt_BR: Egypt + - value: SV + label: + en_US: El Salvador + zh_Hans: 萨尔瓦多 + pt_BR: El Salvador + - value: ET + label: + en_US: Ethiopia + zh_Hans: 埃塞俄比亚 + pt_BR: Ethiopia + - value: FK + label: + en_US: Falkland Islands (Malvinas) + zh_Hans: 福克兰群岛(马尔维纳斯) + pt_BR: Falkland Islands (Malvinas) + - value: FR + label: + en_US: France + zh_Hans: 法国 + pt_BR: France + - value: GF + label: + en_US: French Guiana + zh_Hans: 法属圭亚那 + pt_BR: French Guiana + - value: PF + label: + en_US: French Polynesia + zh_Hans: 法属波利尼西亚 + pt_BR: French Polynesia + - value: TF + label: + en_US: French Southern Territories + zh_Hans: 法属南部领地 + pt_BR: French Southern Territories + - value: GA + label: + en_US: Gabon + zh_Hans: 加蓬 + pt_BR: Gabon + - value: GM + label: + en_US: Gambia + zh_Hans: 冈比亚 + pt_BR: Gambia + - value: DE + label: + en_US: Germany + zh_Hans: 德国 + pt_BR: Germany + - value: GH + label: + en_US: Ghana + zh_Hans: 加纳 + pt_BR: Ghana + - value: GR + label: + en_US: Greece + zh_Hans: 希腊 + pt_BR: Greece + - value: GP + label: + en_US: Guadeloupe + zh_Hans: 瓜德罗普 + pt_BR: Guadeloupe + - value: GT + label: + en_US: Guatemala + zh_Hans: 危地马拉 + pt_BR: Guatemala + - value: GY + label: + en_US: Guyana + zh_Hans: 圭亚那 + pt_BR: Guyana + - value: HT + label: + en_US: Haiti + zh_Hans: 海地 + pt_BR: Haiti + - value: HN + label: + en_US: Honduras + zh_Hans: 洪都拉斯 + pt_BR: Honduras + - value: HK + label: + en_US: Hong Kong + zh_Hans: 香港 + pt_BR: Hong Kong + - value: IN + label: + en_US: India + zh_Hans: 印度 + pt_BR: India + - value: ID + label: + en_US: Indonesia + zh_Hans: 印度尼西亚 + pt_BR: Indonesia + - value: IQ + label: + en_US: Iraq + zh_Hans: 伊拉克 + pt_BR: Iraq + - value: IT + label: + en_US: Italy + zh_Hans: 意大利 + pt_BR: Italy + - value: JM + label: + en_US: Jamaica + zh_Hans: 牙买加 + pt_BR: Jamaica + - value: JP + label: + en_US: Japan + zh_Hans: 日本 + pt_BR: Japan + - value: JO + label: + en_US: Jordan + zh_Hans: 约旦 + pt_BR: Jordan + - value: KZ + label: + en_US: Kazakhstan + zh_Hans: 哈萨克斯坦 + pt_BR: Kazakhstan + - value: KE + label: + en_US: Kenya + zh_Hans: 肯尼亚 + pt_BR: Kenya + - value: KW + label: + en_US: Kuwait + zh_Hans: 科威特 + pt_BR: Kuwait + - value: KG + label: + en_US: Kyrgyzstan + zh_Hans: 吉尔吉斯斯坦 + pt_BR: Kyrgyzstan + - value: LB + label: + en_US: Lebanon + zh_Hans: 黎巴嫩 + pt_BR: Lebanon + - value: LS + label: + en_US: Lesotho + zh_Hans: 莱索托 + pt_BR: Lesotho + - value: LY + label: + en_US: Libyan Arab Jamahiriya + zh_Hans: 利比亚 + pt_BR: Libyan Arab Jamahiriya + - value: MG + label: + en_US: Madagascar + zh_Hans: 马达加斯加 + pt_BR: Madagascar + - value: MW + label: + en_US: Malawi + zh_Hans: 马拉维 + pt_BR: Malawi + - value: MY + label: + en_US: Malaysia + zh_Hans: 马来西亚 + pt_BR: Malaysia + - value: ML + label: + en_US: Mali + zh_Hans: 马里 + pt_BR: Mali + - value: MQ + label: + en_US: Martinique + zh_Hans: 马提尼克 + pt_BR: Martinique + - value: MU + label: + en_US: Mauritius + zh_Hans: 毛里求斯 + pt_BR: Mauritius + - value: YT + label: + en_US: Mayotte + zh_Hans: 马约特 + pt_BR: Mayotte + - value: MX + label: + en_US: Mexico + zh_Hans: 墨西哥 + pt_BR: Mexico + - value: MS + label: + en_US: Montserrat + zh_Hans: 蒙特塞拉特 + pt_BR: Montserrat + - value: MA + label: + en_US: Morocco + zh_Hans: 摩洛哥 + pt_BR: Morocco + - value: MZ + label: + en_US: Mozambique + zh_Hans: 莫桑比克 + pt_BR: Mozambique + - value: NA + label: + en_US: Namibia + zh_Hans: 纳米比亚 + pt_BR: Namibia + - value: NL + label: + en_US: Netherlands + zh_Hans: 荷兰 + pt_BR: Netherlands + - value: NC + label: + en_US: New Caledonia + zh_Hans: 新喀里多尼亚 + pt_BR: New Caledonia + - value: NI + label: + en_US: Nicaragua + zh_Hans: 尼加拉瓜 + pt_BR: Nicaragua + - value: NE + label: + en_US: Niger + zh_Hans: 尼日尔 + pt_BR: Niger + - value: NG + label: + en_US: Nigeria + zh_Hans: 尼日利亚 + pt_BR: Nigeria + - value: OM + label: + en_US: Oman + zh_Hans: 阿曼 + pt_BR: Oman + - value: PK + label: + en_US: Pakistan + zh_Hans: 巴基斯坦 + pt_BR: Pakistan + - value: PS + label: + en_US: Palestinian Territory, Occupied + zh_Hans: 巴勒斯坦领土 + pt_BR: Palestinian Territory, Occupied + - value: PA + label: + en_US: Panama + zh_Hans: 巴拿马 + pt_BR: Panama + - value: PY + label: + en_US: Paraguay + zh_Hans: 巴拉圭 + pt_BR: Paraguay + - value: PE + label: + en_US: Peru + zh_Hans: 秘鲁 + pt_BR: Peru + - value: PH + label: + en_US: Philippines + zh_Hans: 菲律宾 + pt_BR: Philippines + - value: PT + label: + en_US: Portugal + zh_Hans: 葡萄牙 + pt_BR: Portugal + - value: PR + label: + en_US: Puerto Rico + zh_Hans: 波多黎各 + pt_BR: Puerto Rico + - value: QA + label: + en_US: Qatar + zh_Hans: 卡塔尔 + pt_BR: Qatar + - value: RE + label: + en_US: Reunion + zh_Hans: 留尼旺 + pt_BR: Reunion + - value: RU + label: + en_US: Russian Federation + zh_Hans: 俄罗斯联邦 + pt_BR: Russian Federation + - value: RW + label: + en_US: Rwanda + zh_Hans: 卢旺达 + pt_BR: Rwanda + - value: SH + label: + en_US: Saint Helena + zh_Hans: 圣赫勒拿 + pt_BR: Saint Helena + - value: PM + label: + en_US: Saint Pierre and Miquelon + zh_Hans: 圣皮埃尔和密克隆 + pt_BR: Saint Pierre and Miquelon + - value: VC + label: + en_US: Saint Vincent and the Grenadines + zh_Hans: 圣文森特和格林纳丁斯 + pt_BR: Saint Vincent and the Grenadines + - value: ST + label: + en_US: Sao Tome and Principe + zh_Hans: 圣多美和普林西比 + pt_BR: Sao Tome and Principe + - value: SA + label: + en_US: Saudi Arabia + zh_Hans: 沙特阿拉伯 + pt_BR: Saudi Arabia + - value: SN + label: + en_US: Senegal + zh_Hans: 塞内加尔 + pt_BR: Senegal + - value: SC + label: + en_US: Seychelles + zh_Hans: 塞舌尔 + pt_BR: Seychelles + - value: SL + label: + en_US: Sierra Leone + zh_Hans: 塞拉利昂 + pt_BR: Sierra Leone + - value: SG + label: + en_US: Singapore + zh_Hans: 新加坡 + pt_BR: Singapore + - value: SO + label: + en_US: Somalia + zh_Hans: 索马里 + pt_BR: Somalia + - value: ZA + label: + en_US: South Africa + zh_Hans: 南非 + pt_BR: South Africa + - value: GS + label: + en_US: South Georgia and the South Sandwich Islands + zh_Hans: 南乔治亚和南桑威奇群岛 + pt_BR: South Georgia and the South Sandwich Islands + - value: ES + label: + en_US: Spain + zh_Hans: 西班牙 + pt_BR: Spain + - value: LK + label: + en_US: Sri Lanka + zh_Hans: 斯里兰卡 + pt_BR: Sri Lanka + - value: SR + label: + en_US: Suriname + zh_Hans: 苏里南 + pt_BR: Suriname + - value: CH + label: + en_US: Switzerland + zh_Hans: 瑞士 + pt_BR: Switzerland + - value: TW + label: + en_US: Taiwan, Province of China + zh_Hans: 中国台湾省 + pt_BR: Taiwan, Province of China + - value: TZ + label: + en_US: Tanzania, United Republic of + zh_Hans: 坦桑尼亚联合共和国 + pt_BR: Tanzania, United Republic of + - value: TH + label: + en_US: Thailand + zh_Hans: 泰国 + pt_BR: Thailand + - value: TG + label: + en_US: Togo + zh_Hans: 多哥 + pt_BR: Togo + - value: TT + label: + en_US: Trinidad and Tobago + zh_Hans: 特立尼达和多巴哥 + pt_BR: Trinidad and Tobago + - value: TN + label: + en_US: Tunisia + zh_Hans: 突尼斯 + pt_BR: Tunisia + - value: TC + label: + en_US: Turks and Caicos Islands + zh_Hans: 特克斯和凯科斯群岛 + pt_BR: Turks and Caicos Islands + - value: UG + label: + en_US: Uganda + zh_Hans: 乌干达 + pt_BR: Uganda + - value: AE + label: + en_US: United Arab Emirates + zh_Hans: 阿联酋 + pt_BR: United Arab Emirates + - value: UK + label: + en_US: United Kingdom + zh_Hans: 英国 + pt_BR: United Kingdom + - value: GB + label: + en_US: United Kingdom + zh_Hans: 英国 + pt_BR: United Kingdom + - value: US + label: + en_US: United States + zh_Hans: 美国 + pt_BR: United States + - value: UY + label: + en_US: Uruguay + zh_Hans: 乌拉圭 + pt_BR: Uruguay + - value: UZ + label: + en_US: Uzbekistan + zh_Hans: 乌兹别克斯坦 + pt_BR: Uzbekistan + - value: VE + label: + en_US: Venezuela + zh_Hans: 委内瑞拉 + pt_BR: Venezuela + - value: VN + label: + en_US: Viet Nam + zh_Hans: 越南 + pt_BR: Viet Nam + - value: VG + label: + en_US: Virgin Islands, British + zh_Hans: 英属维尔京群岛 + pt_BR: Virgin Islands, British + - value: VI + label: + en_US: Virgin Islands, U.S. + zh_Hans: 美属维尔京群岛 + pt_BR: Virgin Islands, U.S. + - value: ZM + label: + en_US: Zambia + zh_Hans: 赞比亚 + pt_BR: Zambia + - value: ZW + label: + en_US: Zimbabwe + zh_Hans: 津巴布韦 + pt_BR: Zimbabwe + - name: hl + type: select + label: + en_US: Language + zh_Hans: 语言 + human_description: + en_US: Defines the interface language of the search. Default is "en". + zh_Hans: 定义搜索的界面语言。默认为“en”。 + required: false + default: en + form: form + options: + - value: af + label: + en_US: Afrikaans + zh_Hans: 南非语 + - value: ak + label: + en_US: Akan + zh_Hans: 阿坎语 + - value: sq + label: + en_US: Albanian + zh_Hans: 阿尔巴尼亚语 + - value: ws + label: + en_US: Samoa + zh_Hans: 萨摩亚语 + - value: am + label: + en_US: Amharic + zh_Hans: 阿姆哈拉语 + - value: ar + label: + en_US: Arabic + zh_Hans: 阿拉伯语 + - value: hy + label: + en_US: Armenian + zh_Hans: 亚美尼亚语 + - value: az + label: + en_US: Azerbaijani + zh_Hans: 阿塞拜疆语 + - value: eu + label: + en_US: Basque + zh_Hans: 巴斯克语 + - value: be + label: + en_US: Belarusian + zh_Hans: 白俄罗斯语 + - value: bem + label: + en_US: Bemba + zh_Hans: 班巴语 + - value: bn + label: + en_US: Bengali + zh_Hans: 孟加拉语 + - value: bh + label: + en_US: Bihari + zh_Hans: 比哈尔语 + - value: xx-bork + label: + en_US: Bork, bork, bork! + zh_Hans: 博克语 + - value: bs + label: + en_US: Bosnian + zh_Hans: 波斯尼亚语 + - value: br + label: + en_US: Breton + zh_Hans: 布列塔尼语 + - value: bg + label: + en_US: Bulgarian + zh_Hans: 保加利亚语 + - value: bt + label: + en_US: Bhutanese + zh_Hans: 不丹语 + - value: km + label: + en_US: Cambodian + zh_Hans: 高棉语 + - value: ca + label: + en_US: Catalan + zh_Hans: 加泰罗尼亚语 + - value: chr + label: + en_US: Cherokee + zh_Hans: 切罗基语 + - value: ny + label: + en_US: Chichewa + zh_Hans: 齐切瓦语 + - value: zh-cn + label: + en_US: Chinese (Simplified) + zh_Hans: 中文(简体) + - value: zh-tw + label: + en_US: Chinese (Traditional) + zh_Hans: 中文(繁体) + - value: co + label: + en_US: Corsican + zh_Hans: 科西嘉语 + - value: hr + label: + en_US: Croatian + zh_Hans: 克罗地亚语 + - value: cs + label: + en_US: Czech + zh_Hans: 捷克语 + - value: da + label: + en_US: Danish + zh_Hans: 丹麦语 + - value: nl + label: + en_US: Dutch + zh_Hans: 荷兰语 + - value: xx-elmer + label: + en_US: Elmer Fudd + zh_Hans: 艾尔默福德语 + - value: en + label: + en_US: English + zh_Hans: 英语 + - value: eo + label: + en_US: Esperanto + zh_Hans: 世界语 + - value: et + label: + en_US: Estonian + zh_Hans: 爱沙尼亚语 + - value: ee + label: + en_US: Ewe + zh_Hans: 埃维语 + - value: fo + label: + en_US: Faroese + zh_Hans: 法罗语 + - value: tl + label: + en_US: Filipino + zh_Hans: 菲律宾语 + - value: fi + label: + en_US: Finnish + zh_Hans: 芬兰语 + - value: fr + label: + en_US: French + zh_Hans: 法语 + - value: fy + label: + en_US: Frisian + zh_Hans: 弗里西亚语 + - value: gaa + label: + en_US: Ga + zh_Hans: 加语 + - value: gl + label: + en_US: Galician + zh_Hans: 加利西亚语 + - value: ka + label: + en_US: Georgian + zh_Hans: 格鲁吉亚语 + - value: de + label: + en_US: German + zh_Hans: 德语 + - value: el + label: + en_US: Greek + zh_Hans: 希腊语 + - value: kl + label: + en_US: Greenlandic + zh_Hans: 格陵兰语 + - value: gn + label: + en_US: Guarani + zh_Hans: 瓜拉尼语 + - value: gu + label: + en_US: Gujarati + zh_Hans: 古吉拉特语 + - value: xx-hacker + label: + en_US: Hacker + zh_Hans: 黑客语 + - value: ht + label: + en_US: Haitian Creole + zh_Hans: 海地克里奥尔语 + - value: ha + label: + en_US: Hausa + zh_Hans: 豪萨语 + - value: haw + label: + en_US: Hawaiian + zh_Hans: 夏威夷语 + - value: iw + label: + en_US: Hebrew + zh_Hans: 希伯来语 + - value: hi + label: + en_US: Hindi + zh_Hans: 印地语 + - value: hu + label: + en_US: Hungarian + zh_Hans: 匈牙利语 + - value: is + label: + en_US: Icelandic + zh_Hans: 冰岛语 + - value: ig + label: + en_US: Igbo + zh_Hans: 伊博语 + - value: id + label: + en_US: Indonesian + zh_Hans: 印尼语 + - value: ia + label: + en_US: Interlingua + zh_Hans: 国际语 + - value: ga + label: + en_US: Irish + zh_Hans: 爱尔兰语 + - value: it + label: + en_US: Italian + zh_Hans: 意大利语 + - value: ja + label: + en_US: Japanese + zh_Hans: 日语 + - value: jw + label: + en_US: Javanese + zh_Hans: 爪哇语 + - value: kn + label: + en_US: Kannada + zh_Hans: 卡纳达语 + - value: kk + label: + en_US: Kazakh + zh_Hans: 哈萨克语 + - value: rw + label: + en_US: Kinyarwanda + zh_Hans: 基尼亚卢旺达语 + - value: rn + label: + en_US: Kirundi + zh_Hans: 基隆迪语 + - value: xx-klingon + label: + en_US: Klingon + zh_Hans: 克林贡语 + - value: kg + label: + en_US: Kongo + zh_Hans: 刚果语 + - value: ko + label: + en_US: Korean + zh_Hans: 韩语 + - value: kri + label: + en_US: Krio (Sierra Leone) + zh_Hans: 塞拉利昂克里奥尔语 + - value: ku + label: + en_US: Kurdish + zh_Hans: 库尔德语 + - value: ckb + label: + en_US: Kurdish (Soranî) + zh_Hans: 库尔德语(索拉尼) + - value: ky + label: + en_US: Kyrgyz + zh_Hans: 吉尔吉斯语 + - value: lo + label: + en_US: Laothian + zh_Hans: 老挝语 + - value: la + label: + en_US: Latin + zh_Hans: 拉丁语 + - value: lv + label: + en_US: Latvian + zh_Hans: 拉脱维亚语 + - value: ln + label: + en_US: Lingala + zh_Hans: 林加拉语 + - value: lt + label: + en_US: Lithuanian + zh_Hans: 立陶宛语 + - value: loz + label: + en_US: Lozi + zh_Hans: 洛齐语 + - value: lg + label: + en_US: Luganda + zh_Hans: 卢干达语 + - value: ach + label: + en_US: Luo + zh_Hans: 卢奥语 + - value: mk + label: + en_US: Macedonian + zh_Hans: 马其顿语 + - value: mg + label: + en_US: Malagasy + zh_Hans: 马尔加什语 + - value: my + label: + en_US: Malay + zh_Hans: 马来语 + - value: ml + label: + en_US: Malayalam + zh_Hans: 马拉雅拉姆语 + - value: mt + label: + en_US: Maltese + zh_Hans: 马耳他语 + - value: mv + label: + en_US: Maldives + zh_Hans: 马尔代夫语 + - value: mi + label: + en_US: Maori + zh_Hans: 毛利语 + - value: mr + label: + en_US: Marathi + zh_Hans: 马拉地语 + - value: mfe + label: + en_US: Mauritian Creole + zh_Hans: 毛里求斯克里奥尔语 + - value: mo + label: + en_US: Moldavian + zh_Hans: 摩尔达维亚语 + - value: mn + label: + en_US: Mongolian + zh_Hans: 蒙古语 + - value: sr-me + label: + en_US: Montenegrin + zh_Hans: 黑山语 + - value: ne + label: + en_US: Nepali + zh_Hans: 尼泊尔语 + - value: pcm + label: + en_US: Nigerian Pidgin + zh_Hans: 尼日利亚皮钦语 + - value: nso + label: + en_US: Northern Sotho + zh_Hans: 北索托语 + - value: "no" + label: + en_US: Norwegian + zh_Hans: 挪威语 + - value: nn + label: + en_US: Norwegian (Nynorsk) + zh_Hans: 挪威语(尼诺斯克语) + - value: oc + label: + en_US: Occitan + zh_Hans: 奥克语 + - value: or + label: + en_US: Oriya + zh_Hans: 奥里亚语 + - value: om + label: + en_US: Oromo + zh_Hans: 奥罗莫语 + - value: ps + label: + en_US: Pashto + zh_Hans: 普什图语 + - value: fa + label: + en_US: Persian + zh_Hans: 波斯语 + - value: xx-pirate + label: + en_US: Pirate + zh_Hans: 海盗语 + - value: pl + label: + en_US: Polish + zh_Hans: 波兰语 + - value: pt + label: + en_US: Portuguese + zh_Hans: 葡萄牙语 + - value: pt-br + label: + en_US: Portuguese (Brazil) + zh_Hans: 葡萄牙语(巴西) + - value: pt-pt + label: + en_US: Portuguese (Portugal) + zh_Hans: 葡萄牙语(葡萄牙) + - value: pa + label: + en_US: Punjabi + zh_Hans: 旁遮普语 + - value: qu + label: + en_US: Quechua + zh_Hans: 克丘亚语 + - value: ro + label: + en_US: Romanian + zh_Hans: 罗马尼亚语 + - value: rm + label: + en_US: Romansh + zh_Hans: 罗曼什语 + - value: nyn + label: + en_US: Runyakitara + zh_Hans: 卢尼亚基塔拉语 + - value: ru + label: + en_US: Russian + zh_Hans: 俄语 + - value: gd + label: + en_US: Scots Gaelic + zh_Hans: 苏格兰盖尔语 + - value: sr + label: + en_US: Serbian + zh_Hans: 塞尔维亚语 + - value: sh + label: + en_US: Serbo-Croatian + zh_Hans: 塞尔维亚-克罗地亚语 + - value: st + label: + en_US: Sesotho + zh_Hans: 塞索托语 + - value: tn + label: + en_US: Setswana + zh_Hans: 塞茨瓦纳语 + - value: crs + label: + en_US: Seychellois Creole + zh_Hans: 塞舌尔克里奥尔语 + - value: sn + label: + en_US: Shona + zh_Hans: 绍纳语 + - value: sd + label: + en_US: Sindhi + zh_Hans: 信德语 + - value: si + label: + en_US: Sinhalese + zh_Hans: 僧伽罗语 + - value: sk + label: + en_US: Slovak + zh_Hans: 斯洛伐克语 + - value: sl + label: + en_US: Slovenian + zh_Hans: 斯洛文尼亚语 + - value: so + label: + en_US: Somali + zh_Hans: 索马里语 + - value: es + label: + en_US: Spanish + zh_Hans: 西班牙语 + - value: es-419 + label: + en_US: Spanish (Latin American) + zh_Hans: 西班牙语(拉丁美洲) + - value: su + label: + en_US: Sundanese + zh_Hans: 巽他语 + - value: sw + label: + en_US: Swahili + zh_Hans: 斯瓦希里语 + - value: sv + label: + en_US: Swedish + zh_Hans: 瑞典语 + - value: tg + label: + en_US: Tajik + zh_Hans: 塔吉克语 + - value: ta + label: + en_US: Tamil + zh_Hans: 泰米尔语 + - value: tt + label: + en_US: Tatar + zh_Hans: 鞑靼语 + - value: te + label: + en_US: Telugu + zh_Hans: 泰卢固语 + - value: th + label: + en_US: Thai + zh_Hans: 泰语 + - value: ti + label: + en_US: Tigrinya + zh_Hans: 提格利尼亚语 + - value: to + label: + en_US: Tonga + zh_Hans: 汤加语 + - value: lua + label: + en_US: Tshiluba + zh_Hans: 卢巴语 + - value: tum + label: + en_US: Tumbuka + zh_Hans: 图布卡语 + - value: tr + label: + en_US: Turkish + zh_Hans: 土耳其语 + - value: tk + label: + en_US: Turkmen + zh_Hans: 土库曼语 + - value: tw + label: + en_US: Twi + zh_Hans: 契维语 + - value: ug + label: + en_US: Uighur + zh_Hans: 维吾尔语 + - value: uk + label: + en_US: Ukrainian + zh_Hans: 乌克兰语 + - value: ur + label: + en_US: Urdu + zh_Hans: 乌尔都语 + - value: uz + label: + en_US: Uzbek + zh_Hans: 乌兹别克语 + - value: vu + label: + en_US: Vanuatu + zh_Hans: 瓦努阿图语 + - value: vi + label: + en_US: Vietnamese + zh_Hans: 越南语 + - value: cy + label: + en_US: Welsh + zh_Hans: 威尔士语 + - value: wo + label: + en_US: Wolof + zh_Hans: 沃洛夫语 + - value: xh + label: + en_US: Xhosa + zh_Hans: 科萨语 + - value: yi + label: + en_US: Yiddish + zh_Hans: 意第绪语 + - value: yo + label: + en_US: Yoruba + zh_Hans: 约鲁巴语 + - value: zu + label: + en_US: Zulu + zh_Hans: 祖鲁语 + - name: is_remote + type: select + label: + en_US: is_remote + zh_Hans: 很遥远 + human_description: + en_US: Filter results based on the work arrangement. Set it to true to find jobs that offer work from home or remote work opportunities. + zh_Hans: 根据工作安排过滤结果。将其设置为 true 可查找提供在家工作或远程工作机会的工作。 + required: false + form: form + options: + - value: 'true' + label: + en_US: "true" + zh_Hans: "true" + - value: 'false' + label: + en_US: "false" + zh_Hans: "false" diff --git a/api/core/tools/provider/builtin/searchapi/tools/google_news.py b/api/core/tools/provider/builtin/searchapi/tools/google_news.py new file mode 100644 index 0000000000000000000000000000000000000000..c8b3ccda05e195a1fabed00bc5c7661279d61cfb --- /dev/null +++ b/api/core/tools/provider/builtin/searchapi/tools/google_news.py @@ -0,0 +1,97 @@ +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +SEARCH_API_URL = "https://www.searchapi.io/api/v1/search" + + +class SearchAPI: + """ + SearchAPI tool provider. + """ + + def __init__(self, api_key: str) -> None: + """Initialize SearchAPI tool provider.""" + self.searchapi_api_key = api_key + + def run(self, query: str, **kwargs: Any) -> str: + """Run query through SearchAPI and parse result.""" + type = kwargs.get("result_type", "text") + return self._process_response(self.results(query, **kwargs), type=type) + + def results(self, query: str, **kwargs: Any) -> dict: + """Run query through SearchAPI and return the raw result.""" + params = self.get_params(query, **kwargs) + response = requests.get( + url=SEARCH_API_URL, + params=params, + headers={"Authorization": f"Bearer {self.searchapi_api_key}"}, + ) + response.raise_for_status() + return response.json() + + def get_params(self, query: str, **kwargs: Any) -> dict[str, str]: + """Get parameters for SearchAPI.""" + return { + "engine": "google_news", + "q": query, + **{key: value for key, value in kwargs.items() if value not in {None, ""}}, + } + + @staticmethod + def _process_response(res: dict, type: str) -> str: + """Process response from SearchAPI.""" + if "error" in res: + return res["error"] + + toret = "" + if type == "text": + if "organic_results" in res and "snippet" in res["organic_results"][0]: + for item in res["organic_results"]: + toret += "content: " + item["snippet"] + "\n" + "link: " + item["link"] + "\n" + if "top_stories" in res and "title" in res["top_stories"][0]: + for item in res["top_stories"]: + toret += "title: " + item["title"] + "\n" + "link: " + item["link"] + "\n" + if toret == "": + toret = "No good search result found" + + elif type == "link": + if "organic_results" in res and "title" in res["organic_results"][0]: + for item in res["organic_results"]: + toret += f"[{item['title']}]({item['link']})\n" + elif "top_stories" in res and "title" in res["top_stories"][0]: + for item in res["top_stories"]: + toret += f"[{item['title']}]({item['link']})\n" + else: + toret = "No good search result found" + return toret + + +class GoogleNewsTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Invoke the SearchApi tool. + """ + query = tool_parameters["query"] + result_type = tool_parameters["result_type"] + num = tool_parameters.get("num", 10) + google_domain = tool_parameters.get("google_domain", "google.com") + gl = tool_parameters.get("gl", "us") + hl = tool_parameters.get("hl", "en") + location = tool_parameters.get("location") + + api_key = self.runtime.credentials["searchapi_api_key"] + result = SearchAPI(api_key).run( + query, result_type=result_type, num=num, google_domain=google_domain, gl=gl, hl=hl, location=location + ) + + if result_type == "text": + return self.create_text_message(text=result) + return self.create_link_message(link=result) diff --git a/api/core/tools/provider/builtin/searchapi/tools/google_news.yaml b/api/core/tools/provider/builtin/searchapi/tools/google_news.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ff34af34cc9f5c5fc29ca97da68836a35e6859d8 --- /dev/null +++ b/api/core/tools/provider/builtin/searchapi/tools/google_news.yaml @@ -0,0 +1,1922 @@ +identity: + name: google_news_api + author: SearchApi + label: + en_US: Google News API + zh_Hans: Google News API +description: + human: + en_US: A tool to retrieve organic search results snippets and links from Google News engine. + zh_Hans: 一种从 Google 新闻引擎检索有机搜索结果片段和链接的工具。 + llm: A tool to retrieve organic search results snippets and links from Google News engine. +parameters: + - name: query + type: string + required: true + label: + en_US: Query + zh_Hans: 询问 + human_description: + en_US: Defines the query you want to search. + zh_Hans: 定义您要搜索的查询。 + llm_description: Defines the search query you want to search. + form: llm + - name: result_type + type: select + required: true + options: + - value: text + label: + en_US: text + zh_Hans: 文本 + - value: link + label: + en_US: link + zh_Hans: 链接 + default: text + label: + en_US: Result type + zh_Hans: 结果类型 + human_description: + en_US: used for selecting the result type, text or link. + zh_Hans: 用于选择结果类型,使用文本还是链接进行展示。 + form: form + - name: location + type: string + required: false + label: + en_US: Location + zh_Hans: 询问 + human_description: + en_US: Defines from where you want the search to originate. (For example - New York) + zh_Hans: 定义您想要搜索的起始位置。 (例如 - 纽约) + llm_description: Defines from where you want the search to originate. (For example - New York) + form: llm + - name: gl + type: select + label: + en_US: Country + zh_Hans: 国家 + required: false + human_description: + en_US: Defines the country of the search. Default is "US". + zh_Hans: 定义搜索的国家/地区。默认为“美国”。 + llm_description: Defines the gl parameter of the Google search. + form: form + default: US + options: + - value: AF + label: + en_US: Afghanistan + zh_Hans: 阿富汗 + pt_BR: Afeganistão + - value: AL + label: + en_US: Albania + zh_Hans: 阿尔巴尼亚 + pt_BR: Albânia + - value: DZ + label: + en_US: Algeria + zh_Hans: 阿尔及利亚 + pt_BR: Argélia + - value: AS + label: + en_US: American Samoa + zh_Hans: 美属萨摩亚 + pt_BR: Samoa Americana + - value: AD + label: + en_US: Andorra + zh_Hans: 安道尔 + pt_BR: Andorra + - value: AO + label: + en_US: Angola + zh_Hans: 安哥拉 + pt_BR: Angola + - value: AI + label: + en_US: Anguilla + zh_Hans: 安圭拉 + pt_BR: Anguilla + - value: AQ + label: + en_US: Antarctica + zh_Hans: 南极洲 + pt_BR: Antártica + - value: AG + label: + en_US: Antigua and Barbuda + zh_Hans: 安提瓜和巴布达 + pt_BR: Antígua e Barbuda + - value: AR + label: + en_US: Argentina + zh_Hans: 阿根廷 + pt_BR: Argentina + - value: AM + label: + en_US: Armenia + zh_Hans: 亚美尼亚 + pt_BR: Armênia + - value: AW + label: + en_US: Aruba + zh_Hans: 阿鲁巴 + pt_BR: Aruba + - value: AU + label: + en_US: Australia + zh_Hans: 澳大利亚 + pt_BR: Austrália + - value: AT + label: + en_US: Austria + zh_Hans: 奥地利 + pt_BR: Áustria + - value: AZ + label: + en_US: Azerbaijan + zh_Hans: 阿塞拜疆 + pt_BR: Azerbaijão + - value: BS + label: + en_US: Bahamas + zh_Hans: 巴哈马 + pt_BR: Bahamas + - value: BH + label: + en_US: Bahrain + zh_Hans: 巴林 + pt_BR: Bahrein + - value: BD + label: + en_US: Bangladesh + zh_Hans: 孟加拉国 + pt_BR: Bangladesh + - value: BB + label: + en_US: Barbados + zh_Hans: 巴巴多斯 + pt_BR: Barbados + - value: BY + label: + en_US: Belarus + zh_Hans: 白俄罗斯 + pt_BR: Bielorrússia + - value: BE + label: + en_US: Belgium + zh_Hans: 比利时 + pt_BR: Bélgica + - value: BZ + label: + en_US: Belize + zh_Hans: 伯利兹 + pt_BR: Belize + - value: BJ + label: + en_US: Benin + zh_Hans: 贝宁 + pt_BR: Benim + - value: BM + label: + en_US: Bermuda + zh_Hans: 百慕大 + pt_BR: Bermudas + - value: BT + label: + en_US: Bhutan + zh_Hans: 不丹 + pt_BR: Butão + - value: BO + label: + en_US: Bolivia + zh_Hans: 玻利维亚 + pt_BR: Bolívia + - value: BA + label: + en_US: Bosnia and Herzegovina + zh_Hans: 波斯尼亚和黑塞哥维那 + pt_BR: Bósnia e Herzegovina + - value: BW + label: + en_US: Botswana + zh_Hans: 博茨瓦纳 + pt_BR: Botsuana + - value: BV + label: + en_US: Bouvet Island + zh_Hans: 布韦岛 + pt_BR: Ilha Bouvet + - value: BR + label: + en_US: Brazil + zh_Hans: 巴西 + pt_BR: Brasil + - value: IO + label: + en_US: British Indian Ocean Territory + zh_Hans: 英属印度洋领地 + pt_BR: Território Britânico do Oceano Índico + - value: BN + label: + en_US: Brunei Darussalam + zh_Hans: 文莱 + pt_BR: Brunei Darussalam + - value: BG + label: + en_US: Bulgaria + zh_Hans: 保加利亚 + pt_BR: Bulgária + - value: BF + label: + en_US: Burkina Faso + zh_Hans: 布基纳法索 + pt_BR: Burkina Faso + - value: BI + label: + en_US: Burundi + zh_Hans: 布隆迪 + pt_BR: Burundi + - value: KH + label: + en_US: Cambodia + zh_Hans: 柬埔寨 + pt_BR: Camboja + - value: CM + label: + en_US: Cameroon + zh_Hans: 喀麦隆 + pt_BR: Camarões + - value: CA + label: + en_US: Canada + zh_Hans: 加拿大 + pt_BR: Canadá + - value: CV + label: + en_US: Cape Verde + zh_Hans: 佛得角 + pt_BR: Cabo Verde + - value: KY + label: + en_US: Cayman Islands + zh_Hans: 开曼群岛 + pt_BR: Ilhas Cayman + - value: CF + label: + en_US: Central African Republic + zh_Hans: 中非共和国 + pt_BR: República Centro-Africana + - value: TD + label: + en_US: Chad + zh_Hans: 乍得 + pt_BR: Chade + - value: CL + label: + en_US: Chile + zh_Hans: 智利 + pt_BR: Chile + - value: CN + label: + en_US: China + zh_Hans: 中国 + pt_BR: China + - value: CX + label: + en_US: Christmas Island + zh_Hans: 圣诞岛 + pt_BR: Ilha do Natal + - value: CC + label: + en_US: Cocos (Keeling) Islands + zh_Hans: 科科斯(基林)群岛 + pt_BR: Ilhas Cocos (Keeling) + - value: CO + label: + en_US: Colombia + zh_Hans: 哥伦比亚 + pt_BR: Colômbia + - value: KM + label: + en_US: Comoros + zh_Hans: 科摩罗 + pt_BR: Comores + - value: CG + label: + en_US: Congo + zh_Hans: 刚果 + pt_BR: Congo + - value: CD + label: + en_US: Congo, the Democratic Republic of the + zh_Hans: 刚果民主共和国 + pt_BR: Congo, República Democrática do + - value: CK + label: + en_US: Cook Islands + zh_Hans: 库克群岛 + pt_BR: Ilhas Cook + - value: CR + label: + en_US: Costa Rica + zh_Hans: 哥斯达黎加 + pt_BR: Costa Rica + - value: CI + label: + en_US: Cote D'ivoire + zh_Hans: 科特迪瓦 + pt_BR: Costa do Marfim + - value: HR + label: + en_US: Croatia + zh_Hans: 克罗地亚 + pt_BR: Croácia + - value: CU + label: + en_US: Cuba + zh_Hans: 古巴 + pt_BR: Cuba + - value: CY + label: + en_US: Cyprus + zh_Hans: 塞浦路斯 + pt_BR: Chipre + - value: CZ + label: + en_US: Czech Republic + zh_Hans: 捷克共和国 + pt_BR: República Tcheca + - value: DK + label: + en_US: Denmark + zh_Hans: 丹麦 + pt_BR: Dinamarca + - value: DJ + label: + en_US: Djibouti + zh_Hans: 吉布提 + pt_BR: Djibuti + - value: DM + label: + en_US: Dominica + zh_Hans: 多米尼克 + pt_BR: Dominica + - value: DO + label: + en_US: Dominican Republic + zh_Hans: 多米尼加共和国 + pt_BR: República Dominicana + - value: EC + label: + en_US: Ecuador + zh_Hans: 厄瓜多尔 + pt_BR: Equador + - value: EG + label: + en_US: Egypt + zh_Hans: 埃及 + pt_BR: Egito + - value: SV + label: + en_US: El Salvador + zh_Hans: 萨尔瓦多 + pt_BR: El Salvador + - value: GQ + label: + en_US: Equatorial Guinea + zh_Hans: 赤道几内亚 + pt_BR: Guiné Equatorial + - value: ER + label: + en_US: Eritrea + zh_Hans: 厄立特里亚 + pt_BR: Eritreia + - value: EE + label: + en_US: Estonia + zh_Hans: 爱沙尼亚 + pt_BR: Estônia + - value: ET + label: + en_US: Ethiopia + zh_Hans: 埃塞俄比亚 + pt_BR: Etiópia + - value: FK + label: + en_US: Falkland Islands (Malvinas) + zh_Hans: 福克兰群岛(马尔维纳斯) + pt_BR: Ilhas Falkland (Malvinas) + - value: FO + label: + en_US: Faroe Islands + zh_Hans: 法罗群岛 + pt_BR: Ilhas Faroe + - value: FJ + label: + en_US: Fiji + zh_Hans: 斐济 + pt_BR: Fiji + - value: FI + label: + en_US: Finland + zh_Hans: 芬兰 + pt_BR: Finlândia + - value: FR + label: + en_US: France + zh_Hans: 法国 + pt_BR: França + - value: GF + label: + en_US: French Guiana + zh_Hans: 法属圭亚那 + pt_BR: Guiana Francesa + - value: PF + label: + en_US: French Polynesia + zh_Hans: 法属波利尼西亚 + pt_BR: Polinésia Francesa + - value: TF + label: + en_US: French Southern Territories + zh_Hans: 法属南部领地 + pt_BR: Territórios Franceses do Sul + - value: GA + label: + en_US: Gabon + zh_Hans: 加蓬 + pt_BR: Gabão + - value: GM + label: + en_US: Gambia + zh_Hans: 冈比亚 + pt_BR: Gâmbia + - value: GE + label: + en_US: Georgia + zh_Hans: 格鲁吉亚 + pt_BR: Geórgia + - value: DE + label: + en_US: Germany + zh_Hans: 德国 + pt_BR: Alemanha + - value: GH + label: + en_US: Ghana + zh_Hans: 加纳 + pt_BR: Gana + - value: GI + label: + en_US: Gibraltar + zh_Hans: 直布罗陀 + pt_BR: Gibraltar + - value: GR + label: + en_US: Greece + zh_Hans: 希腊 + pt_BR: Grécia + - value: GL + label: + en_US: Greenland + zh_Hans: 格陵兰 + pt_BR: Groenlândia + - value: GD + label: + en_US: Grenada + zh_Hans: 格林纳达 + pt_BR: Granada + - value: GP + label: + en_US: Guadeloupe + zh_Hans: 瓜德罗普 + pt_BR: Guadalupe + - value: GU + label: + en_US: Guam + zh_Hans: 关岛 + pt_BR: Guam + - value: GT + label: + en_US: Guatemala + zh_Hans: 危地马拉 + pt_BR: Guatemala + - value: GN + label: + en_US: Guinea + zh_Hans: 几内亚 + pt_BR: Guiné + - value: GW + label: + en_US: Guinea-Bissau + zh_Hans: 几内亚比绍 + pt_BR: Guiné-Bissau + - value: GY + label: + en_US: Guyana + zh_Hans: 圭亚那 + pt_BR: Guiana + - value: HT + label: + en_US: Haiti + zh_Hans: 海地 + pt_BR: Haiti + - value: HM + label: + en_US: Heard Island and McDonald Islands + zh_Hans: 赫德岛和麦克唐纳群岛 + pt_BR: Ilha Heard e Ilhas McDonald + - value: VA + label: + en_US: Holy See (Vatican City State) + zh_Hans: 教廷(梵蒂冈城国) + pt_BR: Santa Sé (Estado da Cidade do Vaticano) + - value: HN + label: + en_US: Honduras + zh_Hans: 洪都拉斯 + pt_BR: Honduras + - value: HK + label: + en_US: Hong Kong + zh_Hans: 香港 + pt_BR: Hong Kong + - value: HU + label: + en_US: Hungary + zh_Hans: 匈牙利 + pt_BR: Hungria + - value: IS + label: + en_US: Iceland + zh_Hans: 冰岛 + pt_BR: Islândia + - value: IN + label: + en_US: India + zh_Hans: 印度 + pt_BR: Índia + - value: ID + label: + en_US: Indonesia + zh_Hans: 印度尼西亚 + pt_BR: Indonésia + - value: IR + label: + en_US: Iran, Islamic Republic of + zh_Hans: 伊朗 + pt_BR: Irã + - value: IQ + label: + en_US: Iraq + zh_Hans: 伊拉克 + pt_BR: Iraque + - value: IE + label: + en_US: Ireland + zh_Hans: 爱尔兰 + pt_BR: Irlanda + - value: IL + label: + en_US: Israel + zh_Hans: 以色列 + pt_BR: Israel + - value: IT + label: + en_US: Italy + zh_Hans: 意大利 + pt_BR: Itália + - value: JM + label: + en_US: Jamaica + zh_Hans: 牙买加 + pt_BR: Jamaica + - value: JP + label: + en_US: Japan + zh_Hans: 日本 + pt_BR: Japão + - value: JO + label: + en_US: Jordan + zh_Hans: 约旦 + pt_BR: Jordânia + - value: KZ + label: + en_US: Kazakhstan + zh_Hans: 哈萨克斯坦 + pt_BR: Cazaquistão + - value: KE + label: + en_US: Kenya + zh_Hans: 肯尼亚 + pt_BR: Quênia + - value: KI + label: + en_US: Kiribati + zh_Hans: 基里巴斯 + pt_BR: Kiribati + - value: KP + label: + en_US: Korea, Democratic People's Republic of + zh_Hans: 朝鲜 + pt_BR: Coreia, República Democrática Popular da + - value: KR + label: + en_US: Korea, Republic of + zh_Hans: 韩国 + pt_BR: Coreia, República da + - value: KW + label: + en_US: Kuwait + zh_Hans: 科威特 + pt_BR: Kuwait + - value: KG + label: + en_US: Kyrgyzstan + zh_Hans: 吉尔吉斯斯坦 + pt_BR: Quirguistão + - value: LA + label: + en_US: Lao People's Democratic Republic + zh_Hans: 老挝 + pt_BR: República Democrática Popular do Laos + - value: LV + label: + en_US: Latvia + zh_Hans: 拉脱维亚 + pt_BR: Letônia + - value: LB + label: + en_US: Lebanon + zh_Hans: 黎巴嫩 + pt_BR: Líbano + - value: LS + label: + en_US: Lesotho + zh_Hans: 莱索托 + pt_BR: Lesoto + - value: LR + label: + en_US: Liberia + zh_Hans: 利比里亚 + pt_BR: Libéria + - value: LY + label: + en_US: Libyan Arab Jamahiriya + zh_Hans: 利比亚 + pt_BR: Líbia + - value: LI + label: + en_US: Liechtenstein + zh_Hans: 列支敦士登 + pt_BR: Liechtenstein + - value: LT + label: + en_US: Lithuania + zh_Hans: 立陶宛 + pt_BR: Lituânia + - value: LU + label: + en_US: Luxembourg + zh_Hans: 卢森堡 + pt_BR: Luxemburgo + - value: MO + label: + en_US: Macao + zh_Hans: 澳门 + pt_BR: Macau + - value: MK + label: + en_US: Macedonia, the Former Yugosalv Republic of + zh_Hans: 前南斯拉夫马其顿共和国 + pt_BR: Macedônia, Ex-República Iugoslava da + - value: MG + label: + en_US: Madagascar + zh_Hans: 马达加斯加 + pt_BR: Madagascar + - value: MW + label: + en_US: Malawi + zh_Hans: 马拉维 + pt_BR: Malaui + - value: MY + label: + en_US: Malaysia + zh_Hans: 马来西亚 + pt_BR: Malásia + - value: MV + label: + en_US: Maldives + zh_Hans: 马尔代夫 + pt_BR: Maldivas + - value: ML + label: + en_US: Mali + zh_Hans: 马里 + pt_BR: Mali + - value: MT + label: + en_US: Malta + zh_Hans: 马耳他 + pt_BR: Malta + - value: MH + label: + en_US: Marshall Islands + zh_Hans: 马绍尔群岛 + pt_BR: Ilhas Marshall + - value: MQ + label: + en_US: Martinique + zh_Hans: 马提尼克 + pt_BR: Martinica + - value: MR + label: + en_US: Mauritania + zh_Hans: 毛里塔尼亚 + pt_BR: Mauritânia + - value: MU + label: + en_US: Mauritius + zh_Hans: 毛里求斯 + pt_BR: Maurício + - value: YT + label: + en_US: Mayotte + zh_Hans: 马约特 + pt_BR: Mayotte + - value: MX + label: + en_US: Mexico + zh_Hans: 墨西哥 + pt_BR: México + - value: FM + label: + en_US: Micronesia, Federated States of + zh_Hans: 密克罗尼西亚联邦 + pt_BR: Micronésia, Estados Federados da + - value: MD + label: + en_US: Moldova, Republic of + zh_Hans: 摩尔多瓦共和国 + pt_BR: Moldávia, República da + - value: MC + label: + en_US: Monaco + zh_Hans: 摩纳哥 + pt_BR: Mônaco + - value: MN + label: + en_US: Mongolia + zh_Hans: 蒙古 + pt_BR: Mongólia + - value: MS + label: + en_US: Montserrat + zh_Hans: 蒙特塞拉特 + pt_BR: Montserrat + - value: MA + label: + en_US: Morocco + zh_Hans: 摩洛哥 + pt_BR: Marrocos + - value: MZ + label: + en_US: Mozambique + zh_Hans: 莫桑比克 + pt_BR: Moçambique + - value: MM + label: + en_US: Myanmar + zh_Hans: 缅甸 + pt_BR: Mianmar + - value: NA + label: + en_US: Namibia + zh_Hans: 纳米比亚 + pt_BR: Namíbia + - value: NR + label: + en_US: Nauru + zh_Hans: 瑙鲁 + pt_BR: Nauru + - value: NP + label: + en_US: Nepal + zh_Hans: 尼泊尔 + pt_BR: Nepal + - value: NL + label: + en_US: Netherlands + zh_Hans: 荷兰 + pt_BR: Países Baixos + - value: AN + label: + en_US: Netherlands Antilles + zh_Hans: 荷属安的列斯 + pt_BR: Antilhas Holandesas + - value: NC + label: + en_US: New Caledonia + zh_Hans: 新喀里多尼亚 + pt_BR: Nova Caledônia + - value: NZ + label: + en_US: New Zealand + zh_Hans: 新西兰 + pt_BR: Nova Zelândia + - value: NI + label: + en_US: Nicaragua + zh_Hans: 尼加拉瓜 + pt_BR: Nicarágua + - value: NE + label: + en_US: Niger + zh_Hans: 尼日尔 + pt_BR: Níger + - value: NG + label: + en_US: Nigeria + zh_Hans: 尼日利亚 + pt_BR: Nigéria + - value: NU + label: + en_US: Niue + zh_Hans: 纽埃 + pt_BR: Niue + - value: NF + label: + en_US: Norfolk Island + zh_Hans: 诺福克岛 + pt_BR: Ilha Norfolk + - value: MP + label: + en_US: Northern Mariana Islands + zh_Hans: 北马里亚纳群岛 + pt_BR: Ilhas Marianas do Norte + - value: "NO" + label: + en_US: Norway + zh_Hans: 挪威 + pt_BR: Noruega + - value: OM + label: + en_US: Oman + zh_Hans: 阿曼 + pt_BR: Omã + - value: PK + label: + en_US: Pakistan + zh_Hans: 巴基斯坦 + pt_BR: Paquistão + - value: PW + label: + en_US: Palau + zh_Hans: 帕劳 + pt_BR: Palau + - value: PS + label: + en_US: Palestinian Territory, Occupied + zh_Hans: 巴勒斯坦领土 + pt_BR: Palestina, Território Ocupado + - value: PA + label: + en_US: Panama + zh_Hans: 巴拿马 + pt_BR: Panamá + - value: PG + label: + en_US: Papua New Guinea + zh_Hans: 巴布亚新几内亚 + pt_BR: Papua Nova Guiné + - value: PY + label: + en_US: Paraguay + zh_Hans: 巴拉圭 + pt_BR: Paraguai + - value: PE + label: + en_US: Peru + zh_Hans: 秘鲁 + pt_BR: Peru + - value: PH + label: + en_US: Philippines + zh_Hans: 菲律宾 + pt_BR: Filipinas + - value: PN + label: + en_US: Pitcairn + zh_Hans: 皮特凯恩岛 + pt_BR: Pitcairn + - value: PL + label: + en_US: Poland + zh_Hans: 波兰 + pt_BR: Polônia + - value: PT + label: + en_US: Portugal + zh_Hans: 葡萄牙 + pt_BR: Portugal + - value: PR + label: + en_US: Puerto Rico + zh_Hans: 波多黎各 + pt_BR: Porto Rico + - value: QA + label: + en_US: Qatar + zh_Hans: 卡塔尔 + pt_BR: Catar + - value: RE + label: + en_US: Reunion + zh_Hans: 留尼旺 + pt_BR: Reunião + - value: RO + label: + en_US: Romania + zh_Hans: 罗马尼亚 + pt_BR: Romênia + - value: RU + label: + en_US: Russian Federation + zh_Hans: 俄罗斯联邦 + pt_BR: Rússia + - value: RW + label: + en_US: Rwanda + zh_Hans: 卢旺达 + pt_BR: Ruanda + - value: SH + label: + en_US: Saint Helena + zh_Hans: 圣赫勒拿 + pt_BR: Santa Helena + - value: KN + label: + en_US: Saint Kitts and Nevis + zh_Hans: 圣基茨和尼维斯 + pt_BR: São Cristóvão e Nevis + - value: LC + label: + en_US: Saint Lucia + zh_Hans: 圣卢西亚 + pt_BR: Santa Lúcia + - value: PM + label: + en_US: Saint Pierre and Miquelon + zh_Hans: 圣皮埃尔和密克隆 + pt_BR: São Pedro e Miquelon + - value: VC + label: + en_US: Saint Vincent and the Grenadines + zh_Hans: 圣文森特和格林纳丁斯 + pt_BR: São Vicente e Granadinas + - value: WS + label: + en_US: Samoa + zh_Hans: 萨摩亚 + pt_BR: Samoa + - value: SM + label: + en_US: San Marino + zh_Hans: 圣马力诺 + pt_BR: San Marino + - value: ST + label: + en_US: Sao Tome and Principe + zh_Hans: 圣多美和普林西比 + pt_BR: São Tomé e Príncipe + - value: SA + label: + en_US: Saudi Arabia + zh_Hans: 沙特阿拉伯 + pt_BR: Arábia Saudita + - value: SN + label: + en_US: Senegal + zh_Hans: 塞内加尔 + pt_BR: Senegal + - value: RS + label: + en_US: Serbia and Montenegro + zh_Hans: 塞尔维亚和黑山 + pt_BR: Sérvia e Montenegro + - value: SC + label: + en_US: Seychelles + zh_Hans: 塞舌尔 + pt_BR: Seicheles + - value: SL + label: + en_US: Sierra Leone + zh_Hans: 塞拉利昂 + pt_BR: Serra Leoa + - value: SG + label: + en_US: Singapore + zh_Hans: 新加坡 + pt_BR: Singapura + - value: SK + label: + en_US: Slovakia + zh_Hans: 斯洛伐克 + pt_BR: Eslováquia + - value: SI + label: + en_US: Slovenia + zh_Hans: 斯洛文尼亚 + pt_BR: Eslovênia + - value: SB + label: + en_US: Solomon Islands + zh_Hans: 所罗门群岛 + pt_BR: Ilhas Salomão + - value: SO + label: + en_US: Somalia + zh_Hans: 索马里 + pt_BR: Somália + - value: ZA + label: + en_US: South Africa + zh_Hans: 南非 + pt_BR: África do Sul + - value: GS + label: + en_US: South Georgia and the South Sandwich Islands + zh_Hans: 南乔治亚和南桑威奇群岛 + pt_BR: Geórgia do Sul e Ilhas Sandwich do Sul + - value: ES + label: + en_US: Spain + zh_Hans: 西班牙 + pt_BR: Espanha + - value: LK + label: + en_US: Sri Lanka + zh_Hans: 斯里兰卡 + pt_BR: Sri Lanka + - value: SD + label: + en_US: Sudan + zh_Hans: 苏丹 + pt_BR: Sudão + - value: SR + label: + en_US: Suriname + zh_Hans: 苏里南 + pt_BR: Suriname + - value: SJ + label: + en_US: Svalbard and Jan Mayen + zh_Hans: 斯瓦尔巴特和扬马延岛 + pt_BR: Svalbard e Jan Mayen + - value: SZ + label: + en_US: Swaziland + zh_Hans: 斯威士兰 + pt_BR: Essuatíni + - value: SE + label: + en_US: Sweden + zh_Hans: 瑞典 + pt_BR: Suécia + - value: CH + label: + en_US: Switzerland + zh_Hans: 瑞士 + pt_BR: Suíça + - value: SY + label: + en_US: Syrian Arab Republic + zh_Hans: 叙利亚 + pt_BR: Síria + - value: TW + label: + en_US: Taiwan, Province of China + zh_Hans: 台湾 + pt_BR: Taiwan + - value: TJ + label: + en_US: Tajikistan + zh_Hans: 塔吉克斯坦 + pt_BR: Tajiquistão + - value: TZ + label: + en_US: Tanzania, United Republic of + zh_Hans: 坦桑尼亚联合共和国 + pt_BR: Tanzânia + - value: TH + label: + en_US: Thailand + zh_Hans: 泰国 + pt_BR: Tailândia + - value: TL + label: + en_US: Timor-Leste + zh_Hans: 东帝汶 + pt_BR: Timor-Leste + - value: TG + label: + en_US: Togo + zh_Hans: 多哥 + pt_BR: Togo + - value: TK + label: + en_US: Tokelau + zh_Hans: 托克劳 + pt_BR: Toquelau + - value: TO + label: + en_US: Tonga + zh_Hans: 汤加 + pt_BR: Tonga + - value: TT + label: + en_US: Trinidad and Tobago + zh_Hans: 特立尼达和多巴哥 + pt_BR: Trindade e Tobago + - value: TN + label: + en_US: Tunisia + zh_Hans: 突尼斯 + pt_BR: Tunísia + - value: TR + label: + en_US: Turkey + zh_Hans: 土耳其 + pt_BR: Turquia + - value: TM + label: + en_US: Turkmenistan + zh_Hans: 土库曼斯坦 + pt_BR: Turcomenistão + - value: TC + label: + en_US: Turks and Caicos Islands + zh_Hans: 特克斯和凯科斯群岛 + pt_BR: Ilhas Turks e Caicos + - value: TV + label: + en_US: Tuvalu + zh_Hans: 图瓦卢 + pt_BR: Tuvalu + - value: UG + label: + en_US: Uganda + zh_Hans: 乌干达 + pt_BR: Uganda + - value: UA + label: + en_US: Ukraine + zh_Hans: 乌克兰 + pt_BR: Ucrânia + - value: AE + label: + en_US: United Arab Emirates + zh_Hans: 阿联酋 + pt_BR: Emirados Árabes Unidos + - value: UK + label: + en_US: United Kingdom + zh_Hans: 英国 + pt_BR: Reino Unido + - value: GB + label: + en_US: United Kingdom + zh_Hans: 英国 + pt_BR: Reino Unido + - value: US + label: + en_US: United States + zh_Hans: 美国 + pt_BR: Estados Unidos + - value: UM + label: + en_US: United States Minor Outlying Islands + zh_Hans: 美国本土外小岛屿 + pt_BR: Ilhas Menores Distantes dos Estados Unidos + - value: UY + label: + en_US: Uruguay + zh_Hans: 乌拉圭 + pt_BR: Uruguai + - value: UZ + label: + en_US: Uzbekistan + zh_Hans: 乌兹别克斯坦 + pt_BR: Uzbequistão + - value: VU + label: + en_US: Vanuatu + zh_Hans: 瓦努阿图 + pt_BR: Vanuatu + - value: VE + label: + en_US: Venezuela + zh_Hans: 委内瑞拉 + pt_BR: Venezuela + - value: VN + label: + en_US: Viet Nam + zh_Hans: 越南 + pt_BR: Vietnã + - value: VG + label: + en_US: Virgin Islands, British + zh_Hans: 英属维尔京群岛 + pt_BR: Ilhas Virgens Britânicas + - value: VI + label: + en_US: Virgin Islands, U.S. + zh_Hans: 美属维尔京群岛 + pt_BR: Ilhas Virgens dos EUA + - value: WF + label: + en_US: Wallis and Futuna + zh_Hans: 瓦利斯和富图纳群岛 + pt_BR: Wallis e Futuna + - value: EH + label: + en_US: Western Sahara + zh_Hans: 西撒哈拉 + pt_BR: Saara Ocidental + - value: YE + label: + en_US: Yemen + zh_Hans: 也门 + pt_BR: Iémen + - value: ZM + label: + en_US: Zambia + zh_Hans: 赞比亚 + pt_BR: Zâmbia + - value: ZW + label: + en_US: Zimbabwe + zh_Hans: 津巴布韦 + pt_BR: Zimbábue + - name: hl + type: select + label: + en_US: Language + zh_Hans: 语言 + human_description: + en_US: Defines the interface language of the search. Default is "en". + zh_Hans: 定义搜索的界面语言。默认为“en”。 + required: false + default: en + form: form + options: + - value: af + label: + en_US: Afrikaans + zh_Hans: 南非语 + - value: ak + label: + en_US: Akan + zh_Hans: 阿坎语 + - value: sq + label: + en_US: Albanian + zh_Hans: 阿尔巴尼亚语 + - value: ws + label: + en_US: Samoa + zh_Hans: 萨摩亚语 + - value: am + label: + en_US: Amharic + zh_Hans: 阿姆哈拉语 + - value: ar + label: + en_US: Arabic + zh_Hans: 阿拉伯语 + - value: hy + label: + en_US: Armenian + zh_Hans: 亚美尼亚语 + - value: az + label: + en_US: Azerbaijani + zh_Hans: 阿塞拜疆语 + - value: eu + label: + en_US: Basque + zh_Hans: 巴斯克语 + - value: be + label: + en_US: Belarusian + zh_Hans: 白俄罗斯语 + - value: bem + label: + en_US: Bemba + zh_Hans: 班巴语 + - value: bn + label: + en_US: Bengali + zh_Hans: 孟加拉语 + - value: bh + label: + en_US: Bihari + zh_Hans: 比哈尔语 + - value: xx-bork + label: + en_US: Bork, bork, bork! + zh_Hans: 博克语 + - value: bs + label: + en_US: Bosnian + zh_Hans: 波斯尼亚语 + - value: br + label: + en_US: Breton + zh_Hans: 布列塔尼语 + - value: bg + label: + en_US: Bulgarian + zh_Hans: 保加利亚语 + - value: bt + label: + en_US: Bhutanese + zh_Hans: 不丹语 + - value: km + label: + en_US: Cambodian + zh_Hans: 高棉语 + - value: ca + label: + en_US: Catalan + zh_Hans: 加泰罗尼亚语 + - value: chr + label: + en_US: Cherokee + zh_Hans: 切罗基语 + - value: ny + label: + en_US: Chichewa + zh_Hans: 齐切瓦语 + - value: zh-cn + label: + en_US: Chinese (Simplified) + zh_Hans: 中文(简体) + - value: zh-tw + label: + en_US: Chinese (Traditional) + zh_Hans: 中文(繁体) + - value: co + label: + en_US: Corsican + zh_Hans: 科西嘉语 + - value: hr + label: + en_US: Croatian + zh_Hans: 克罗地亚语 + - value: cs + label: + en_US: Czech + zh_Hans: 捷克语 + - value: da + label: + en_US: Danish + zh_Hans: 丹麦语 + - value: nl + label: + en_US: Dutch + zh_Hans: 荷兰语 + - value: xx-elmer + label: + en_US: Elmer Fudd + zh_Hans: 艾尔默福德语 + - value: en + label: + en_US: English + zh_Hans: 英语 + - value: eo + label: + en_US: Esperanto + zh_Hans: 世界语 + - value: et + label: + en_US: Estonian + zh_Hans: 爱沙尼亚语 + - value: ee + label: + en_US: Ewe + zh_Hans: 埃维语 + - value: fo + label: + en_US: Faroese + zh_Hans: 法罗语 + - value: tl + label: + en_US: Filipino + zh_Hans: 菲律宾语 + - value: fi + label: + en_US: Finnish + zh_Hans: 芬兰语 + - value: fr + label: + en_US: French + zh_Hans: 法语 + - value: fy + label: + en_US: Frisian + zh_Hans: 弗里西亚语 + - value: gaa + label: + en_US: Ga + zh_Hans: 加语 + - value: gl + label: + en_US: Galician + zh_Hans: 加利西亚语 + - value: ka + label: + en_US: Georgian + zh_Hans: 格鲁吉亚语 + - value: de + label: + en_US: German + zh_Hans: 德语 + - value: el + label: + en_US: Greek + zh_Hans: 希腊语 + - value: kl + label: + en_US: Greenlandic + zh_Hans: 格陵兰语 + - value: gn + label: + en_US: Guarani + zh_Hans: 瓜拉尼语 + - value: gu + label: + en_US: Gujarati + zh_Hans: 古吉拉特语 + - value: xx-hacker + label: + en_US: Hacker + zh_Hans: 黑客语 + - value: ht + label: + en_US: Haitian Creole + zh_Hans: 海地克里奥尔语 + - value: ha + label: + en_US: Hausa + zh_Hans: 豪萨语 + - value: haw + label: + en_US: Hawaiian + zh_Hans: 夏威夷语 + - value: iw + label: + en_US: Hebrew + zh_Hans: 希伯来语 + - value: hi + label: + en_US: Hindi + zh_Hans: 印地语 + - value: hu + label: + en_US: Hungarian + zh_Hans: 匈牙利语 + - value: is + label: + en_US: Icelandic + zh_Hans: 冰岛语 + - value: ig + label: + en_US: Igbo + zh_Hans: 伊博语 + - value: id + label: + en_US: Indonesian + zh_Hans: 印尼语 + - value: ia + label: + en_US: Interlingua + zh_Hans: 国际语 + - value: ga + label: + en_US: Irish + zh_Hans: 爱尔兰语 + - value: it + label: + en_US: Italian + zh_Hans: 意大利语 + - value: ja + label: + en_US: Japanese + zh_Hans: 日语 + - value: jw + label: + en_US: Javanese + zh_Hans: 爪哇语 + - value: kn + label: + en_US: Kannada + zh_Hans: 卡纳达语 + - value: kk + label: + en_US: Kazakh + zh_Hans: 哈萨克语 + - value: rw + label: + en_US: Kinyarwanda + zh_Hans: 基尼亚卢旺达语 + - value: rn + label: + en_US: Kirundi + zh_Hans: 基隆迪语 + - value: xx-klingon + label: + en_US: Klingon + zh_Hans: 克林贡语 + - value: kg + label: + en_US: Kongo + zh_Hans: 刚果语 + - value: ko + label: + en_US: Korean + zh_Hans: 韩语 + - value: kri + label: + en_US: Krio (Sierra Leone) + zh_Hans: 塞拉利昂克里奥尔语 + - value: ku + label: + en_US: Kurdish + zh_Hans: 库尔德语 + - value: ckb + label: + en_US: Kurdish (Soranî) + zh_Hans: 库尔德语(索拉尼) + - value: ky + label: + en_US: Kyrgyz + zh_Hans: 吉尔吉斯语 + - value: lo + label: + en_US: Laothian + zh_Hans: 老挝语 + - value: la + label: + en_US: Latin + zh_Hans: 拉丁语 + - value: lv + label: + en_US: Latvian + zh_Hans: 拉脱维亚语 + - value: ln + label: + en_US: Lingala + zh_Hans: 林加拉语 + - value: lt + label: + en_US: Lithuanian + zh_Hans: 立陶宛语 + - value: loz + label: + en_US: Lozi + zh_Hans: 洛齐语 + - value: lg + label: + en_US: Luganda + zh_Hans: 卢干达语 + - value: ach + label: + en_US: Luo + zh_Hans: 卢奥语 + - value: mk + label: + en_US: Macedonian + zh_Hans: 马其顿语 + - value: mg + label: + en_US: Malagasy + zh_Hans: 马尔加什语 + - value: my + label: + en_US: Malay + zh_Hans: 马来语 + - value: ml + label: + en_US: Malayalam + zh_Hans: 马拉雅拉姆语 + - value: mt + label: + en_US: Maltese + zh_Hans: 马耳他语 + - value: mv + label: + en_US: Maldives + zh_Hans: 马尔代夫语 + - value: mi + label: + en_US: Maori + zh_Hans: 毛利语 + - value: mr + label: + en_US: Marathi + zh_Hans: 马拉地语 + - value: mfe + label: + en_US: Mauritian Creole + zh_Hans: 毛里求斯克里奥尔语 + - value: mo + label: + en_US: Moldavian + zh_Hans: 摩尔达维亚语 + - value: mn + label: + en_US: Mongolian + zh_Hans: 蒙古语 + - value: sr-me + label: + en_US: Montenegrin + zh_Hans: 黑山语 + - value: ne + label: + en_US: Nepali + zh_Hans: 尼泊尔语 + - value: pcm + label: + en_US: Nigerian Pidgin + zh_Hans: 尼日利亚皮钦语 + - value: nso + label: + en_US: Northern Sotho + zh_Hans: 北索托语 + - value: "no" + label: + en_US: Norwegian + zh_Hans: 挪威语 + - value: nn + label: + en_US: Norwegian (Nynorsk) + zh_Hans: 挪威语(尼诺斯克语) + - value: oc + label: + en_US: Occitan + zh_Hans: 奥克语 + - value: or + label: + en_US: Oriya + zh_Hans: 奥里亚语 + - value: om + label: + en_US: Oromo + zh_Hans: 奥罗莫语 + - value: ps + label: + en_US: Pashto + zh_Hans: 普什图语 + - value: fa + label: + en_US: Persian + zh_Hans: 波斯语 + - value: xx-pirate + label: + en_US: Pirate + zh_Hans: 海盗语 + - value: pl + label: + en_US: Polish + zh_Hans: 波兰语 + - value: pt + label: + en_US: Portuguese + zh_Hans: 葡萄牙语 + - value: pt-br + label: + en_US: Portuguese (Brazil) + zh_Hans: 葡萄牙语(巴西) + - value: pt-pt + label: + en_US: Portuguese (Portugal) + zh_Hans: 葡萄牙语(葡萄牙) + - value: pa + label: + en_US: Punjabi + zh_Hans: 旁遮普语 + - value: qu + label: + en_US: Quechua + zh_Hans: 克丘亚语 + - value: ro + label: + en_US: Romanian + zh_Hans: 罗马尼亚语 + - value: rm + label: + en_US: Romansh + zh_Hans: 罗曼什语 + - value: nyn + label: + en_US: Runyakitara + zh_Hans: 卢尼亚基塔拉语 + - value: ru + label: + en_US: Russian + zh_Hans: 俄语 + - value: gd + label: + en_US: Scots Gaelic + zh_Hans: 苏格兰盖尔语 + - value: sr + label: + en_US: Serbian + zh_Hans: 塞尔维亚语 + - value: sh + label: + en_US: Serbo-Croatian + zh_Hans: 塞尔维亚-克罗地亚语 + - value: st + label: + en_US: Sesotho + zh_Hans: 塞索托语 + - value: tn + label: + en_US: Setswana + zh_Hans: 塞茨瓦纳语 + - value: crs + label: + en_US: Seychellois Creole + zh_Hans: 塞舌尔克里奥尔语 + - value: sn + label: + en_US: Shona + zh_Hans: 绍纳语 + - value: sd + label: + en_US: Sindhi + zh_Hans: 信德语 + - value: si + label: + en_US: Sinhalese + zh_Hans: 僧伽罗语 + - value: sk + label: + en_US: Slovak + zh_Hans: 斯洛伐克语 + - value: sl + label: + en_US: Slovenian + zh_Hans: 斯洛文尼亚语 + - value: so + label: + en_US: Somali + zh_Hans: 索马里语 + - value: es + label: + en_US: Spanish + zh_Hans: 西班牙语 + - value: es-419 + label: + en_US: Spanish (Latin American) + zh_Hans: 西班牙语(拉丁美洲) + - value: su + label: + en_US: Sundanese + zh_Hans: 巽他语 + - value: sw + label: + en_US: Swahili + zh_Hans: 斯瓦希里语 + - value: sv + label: + en_US: Swedish + zh_Hans: 瑞典语 + - value: tg + label: + en_US: Tajik + zh_Hans: 塔吉克语 + - value: ta + label: + en_US: Tamil + zh_Hans: 泰米尔语 + - value: tt + label: + en_US: Tatar + zh_Hans: 鞑靼语 + - value: te + label: + en_US: Telugu + zh_Hans: 泰卢固语 + - value: th + label: + en_US: Thai + zh_Hans: 泰语 + - value: ti + label: + en_US: Tigrinya + zh_Hans: 提格利尼亚语 + - value: to + label: + en_US: Tonga + zh_Hans: 汤加语 + - value: lua + label: + en_US: Tshiluba + zh_Hans: 卢巴语 + - value: tum + label: + en_US: Tumbuka + zh_Hans: 图布卡语 + - value: tr + label: + en_US: Turkish + zh_Hans: 土耳其语 + - value: tk + label: + en_US: Turkmen + zh_Hans: 土库曼语 + - value: tw + label: + en_US: Twi + zh_Hans: 契维语 + - value: ug + label: + en_US: Uighur + zh_Hans: 维吾尔语 + - value: uk + label: + en_US: Ukrainian + zh_Hans: 乌克兰语 + - value: ur + label: + en_US: Urdu + zh_Hans: 乌尔都语 + - value: uz + label: + en_US: Uzbek + zh_Hans: 乌兹别克语 + - value: vu + label: + en_US: Vanuatu + zh_Hans: 瓦努阿图语 + - value: vi + label: + en_US: Vietnamese + zh_Hans: 越南语 + - value: cy + label: + en_US: Welsh + zh_Hans: 威尔士语 + - value: wo + label: + en_US: Wolof + zh_Hans: 沃洛夫语 + - value: xh + label: + en_US: Xhosa + zh_Hans: 科萨语 + - value: yi + label: + en_US: Yiddish + zh_Hans: 意第绪语 + - value: yo + label: + en_US: Yoruba + zh_Hans: 约鲁巴语 + - value: zu + label: + en_US: Zulu + zh_Hans: 祖鲁语 + - name: google_domain + type: string + required: false + label: + en_US: google_domain + zh_Hans: google_domain + human_description: + en_US: Defines the Google domain of the search. Default is "google.com". + zh_Hans: 定义搜索的 Google 域。默认为“google.com”。 + llm_description: Defines Google domain in which you want to search. + form: llm + - name: num + type: number + required: false + label: + en_US: num + zh_Hans: num + human_description: + en_US: Specifies the number of results to display per page. Default is 10. Max number - 100, min - 1. + zh_Hans: 指定每页显示的结果数。默认值为 10。最大数量 - 100,最小数量 - 1。 + pt_BR: Specifies the number of results to display per page. Default is 10. Max number - 100, min - 1. + llm_description: Specifies the num of results to display per page. + form: llm diff --git a/api/core/tools/provider/builtin/searchapi/tools/youtube_transcripts.py b/api/core/tools/provider/builtin/searchapi/tools/youtube_transcripts.py new file mode 100644 index 0000000000000000000000000000000000000000..b14821f8312dd01fe5fa2aa2c85a199069ae5822 --- /dev/null +++ b/api/core/tools/provider/builtin/searchapi/tools/youtube_transcripts.py @@ -0,0 +1,75 @@ +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +SEARCH_API_URL = "https://www.searchapi.io/api/v1/search" + + +class SearchAPI: + """ + SearchAPI tool provider. + """ + + def __init__(self, api_key: str) -> None: + """Initialize SearchAPI tool provider.""" + self.searchapi_api_key = api_key + + def run(self, video_id: str, language: str, **kwargs: Any) -> str: + """Run video_id through SearchAPI and parse result.""" + return self._process_response(self.results(video_id, language, **kwargs)) + + def results(self, video_id: str, language: str, **kwargs: Any) -> dict: + """Run video_id through SearchAPI and return the raw result.""" + params = self.get_params(video_id, language, **kwargs) + response = requests.get( + url=SEARCH_API_URL, + params=params, + headers={"Authorization": f"Bearer {self.searchapi_api_key}"}, + ) + response.raise_for_status() + return response.json() + + def get_params(self, video_id: str, language: str, **kwargs: Any) -> dict[str, str]: + """Get parameters for SearchAPI.""" + return { + "engine": "youtube_transcripts", + "video_id": video_id, + "lang": language or "en", + **{key: value for key, value in kwargs.items() if value not in {None, ""}}, + } + + @staticmethod + def _process_response(res: dict) -> str: + """Process response from SearchAPI.""" + if "error" in res: + return res["error"] + + toret = "" + if "transcripts" in res and "text" in res["transcripts"][0]: + for item in res["transcripts"]: + toret += item["text"] + " " + if toret == "": + toret = "No good search result found" + + return toret + + +class YoutubeTranscriptsTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Invoke the SearchApi tool. + """ + video_id = tool_parameters["video_id"] + language = tool_parameters.get("language", "en") + + api_key = self.runtime.credentials["searchapi_api_key"] + result = SearchAPI(api_key).run(video_id, language=language) + + return self.create_text_message(text=result) diff --git a/api/core/tools/provider/builtin/searchapi/tools/youtube_transcripts.yaml b/api/core/tools/provider/builtin/searchapi/tools/youtube_transcripts.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8bdcd6bb936d9632c7f897415207458d4ecd4383 --- /dev/null +++ b/api/core/tools/provider/builtin/searchapi/tools/youtube_transcripts.yaml @@ -0,0 +1,34 @@ +identity: + name: youtube_transcripts_api + author: SearchApi + label: + en_US: YouTube Transcripts API + zh_Hans: YouTube 脚本 API +description: + human: + en_US: A tool to retrieve transcripts from the specific YouTube video. + zh_Hans: 一种从特定 YouTube 视频检索文字记录的工具。 + llm: A tool to retrieve transcripts from the specific YouTube video. +parameters: + - name: video_id + type: string + required: true + label: + en_US: video_id + zh_Hans: 视频ID + human_description: + en_US: Used to define the video you want to search. You can find the video id's in YouTube page that appears in URL. For example - https://www.youtube.com/watch?v=video_id. + zh_Hans: 用于定义要搜索的视频。您可以在 URL 中显示的 YouTube 页面中找到视频 ID。例如 - https://www.youtube.com/watch?v=video_id。 + llm_description: Used to define the video you want to search. + form: llm + - name: language + type: string + required: false + label: + en_US: language + zh_Hans: 语言 + human_description: + en_US: Used to set the language for transcripts. The default value is "en". You can find all supported languages in SearchApi documentation. + zh_Hans: 用于设置成绩单的语言。默认值为“en”。您可以在 SearchApi 文档中找到所有支持的语言。 + llm_description: Used to set the language for transcripts. + form: llm diff --git a/api/core/tools/provider/builtin/searxng/_assets/icon.svg b/api/core/tools/provider/builtin/searxng/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..b94fe3728adbff80cfe769935cf072ada373c18c --- /dev/null +++ b/api/core/tools/provider/builtin/searxng/_assets/icon.svg @@ -0,0 +1,56 @@ + + + + + + + image/svg+xml + + + + + + + + + + + + diff --git a/api/core/tools/provider/builtin/searxng/docker/settings.yml b/api/core/tools/provider/builtin/searxng/docker/settings.yml new file mode 100644 index 0000000000000000000000000000000000000000..18e18688002cbcd2934704aadc1fae4e149f7716 --- /dev/null +++ b/api/core/tools/provider/builtin/searxng/docker/settings.yml @@ -0,0 +1,2501 @@ +general: + # Debug mode, only for development. Is overwritten by ${SEARXNG_DEBUG} + debug: false + # displayed name + instance_name: "searxng" + # For example: https://example.com/privacy + privacypolicy_url: false + # use true to use your own donation page written in searx/info/en/donate.md + # use false to disable the donation link + donation_url: false + # mailto:contact@example.com + contact_url: false + # record stats + enable_metrics: true + +brand: + new_issue_url: https://github.com/searxng/searxng/issues/new + docs_url: https://docs.searxng.org/ + public_instances: https://searx.space + wiki_url: https://github.com/searxng/searxng/wiki + issue_url: https://github.com/searxng/searxng/issues + # custom: + # maintainer: "Jon Doe" + # # Custom entries in the footer: [title]: [link] + # links: + # Uptime: https://uptime.searxng.org/history/darmarit-org + # About: "https://searxng.org" + +search: + # Filter results. 0: None, 1: Moderate, 2: Strict + safe_search: 0 + # Existing autocomplete backends: "dbpedia", "duckduckgo", "google", "yandex", "mwmbl", + # "seznam", "startpage", "stract", "swisscows", "qwant", "wikipedia" - leave blank to turn it off + # by default. + autocomplete: "" + # minimun characters to type before autocompleter starts + autocomplete_min: 4 + # Default search language - leave blank to detect from browser information or + # use codes from 'languages.py' + default_lang: "auto" + # max_page: 0 # if engine supports paging, 0 means unlimited numbers of pages + # Available languages + # languages: + # - all + # - en + # - en-US + # - de + # - it-IT + # - fr + # - fr-BE + # ban time in seconds after engine errors + ban_time_on_fail: 5 + # max ban time in seconds after engine errors + max_ban_time_on_fail: 120 + suspended_times: + # Engine suspension time after error (in seconds; set to 0 to disable) + # For error "Access denied" and "HTTP error [402, 403]" + SearxEngineAccessDenied: 86400 + # For error "CAPTCHA" + SearxEngineCaptcha: 86400 + # For error "Too many request" and "HTTP error 429" + SearxEngineTooManyRequests: 3600 + # Cloudflare CAPTCHA + cf_SearxEngineCaptcha: 1296000 + cf_SearxEngineAccessDenied: 86400 + # ReCAPTCHA + recaptcha_SearxEngineCaptcha: 604800 + + # remove format to deny access, use lower case. + # formats: [html, csv, json, rss] + formats: + - html + - json + +server: + # Is overwritten by ${SEARXNG_PORT} and ${SEARXNG_BIND_ADDRESS} + port: 8888 + bind_address: "127.0.0.1" + # public URL of the instance, to ensure correct inbound links. Is overwritten + # by ${SEARXNG_URL}. + base_url: http://0.0.0.0:8081/ # "http://example.com/location" + # rate limit the number of request on the instance, block some bots. + # Is overwritten by ${SEARXNG_LIMITER} + limiter: false + # enable features designed only for public instances. + # Is overwritten by ${SEARXNG_PUBLIC_INSTANCE} + public_instance: false + + # If your instance owns a /etc/searxng/settings.yml file, then set the following + # values there. + + secret_key: "772ba36386fb56d0f8fe818941552dabbe69220d4c0eb4a385a5729cdbc20c2d" # Is overwritten by ${SEARXNG_SECRET} + # Proxy image results through SearXNG. Is overwritten by ${SEARXNG_IMAGE_PROXY} + image_proxy: false + # 1.0 and 1.1 are supported + http_protocol_version: "1.0" + # POST queries are more secure as they don't show up in history but may cause + # problems when using Firefox containers + method: "POST" + default_http_headers: + X-Content-Type-Options: nosniff + X-Download-Options: noopen + X-Robots-Tag: noindex, nofollow + Referrer-Policy: no-referrer + +redis: + # URL to connect redis database. Is overwritten by ${SEARXNG_REDIS_URL}. + # https://docs.searxng.org/admin/settings/settings_redis.html#settings-redis + url: false + +ui: + # Custom static path - leave it blank if you didn't change + static_path: "" + # Is overwritten by ${SEARXNG_STATIC_USE_HASH}. + static_use_hash: false + # Custom templates path - leave it blank if you didn't change + templates_path: "" + # query_in_title: When true, the result page's titles contains the query + # it decreases the privacy, since the browser can records the page titles. + query_in_title: false + # infinite_scroll: When true, automatically loads the next page when scrolling to bottom of the current page. + infinite_scroll: false + # ui theme + default_theme: simple + # center the results ? + center_alignment: false + # URL prefix of the internet archive, don't forget trailing slash (if needed). + # cache_url: "https://webcache.googleusercontent.com/search?q=cache:" + # Default interface locale - leave blank to detect from browser information or + # use codes from the 'locales' config section + default_locale: "" + # Open result links in a new tab by default + # results_on_new_tab: false + theme_args: + # style of simple theme: auto, light, dark + simple_style: auto + # Perform search immediately if a category selected. + # Disable to select multiple categories at once and start the search manually. + search_on_category_select: true + # Hotkeys: default or vim + hotkeys: default + +# Lock arbitrary settings on the preferences page. To find the ID of the user +# setting you want to lock, check the ID of the form on the page "preferences". +# +# preferences: +# lock: +# - language +# - autocomplete +# - method +# - query_in_title + +# searx supports result proxification using an external service: +# https://github.com/asciimoo/morty uncomment below section if you have running +# morty proxy the key is base64 encoded (keep the !!binary notation) +# Note: since commit af77ec3, morty accepts a base64 encoded key. +# +# result_proxy: +# url: http://127.0.0.1:3000/ +# # the key is a base64 encoded string, the YAML !!binary prefix is optional +# key: !!binary "your_morty_proxy_key" +# # [true|false] enable the "proxy" button next to each result +# proxify_results: true + +# communication with search engines +# +outgoing: + # default timeout in seconds, can be override by engine + request_timeout: 3.0 + # the maximum timeout in seconds + # max_request_timeout: 10.0 + # suffix of searx_useragent, could contain information like an email address + # to the administrator + useragent_suffix: "" + # The maximum number of concurrent connections that may be established. + pool_connections: 100 + # Allow the connection pool to maintain keep-alive connections below this + # point. + pool_maxsize: 20 + # See https://www.python-httpx.org/http2/ + enable_http2: true + # uncomment below section if you want to use a custom server certificate + # see https://www.python-httpx.org/advanced/#changing-the-verification-defaults + # and https://www.python-httpx.org/compatibility/#ssl-configuration + # verify: ~/.mitmproxy/mitmproxy-ca-cert.cer + # + # uncomment below section if you want to use a proxyq see: SOCKS proxies + # https://2.python-requests.org/en/latest/user/advanced/#proxies + # are also supported: see + # https://2.python-requests.org/en/latest/user/advanced/#socks + # + # proxies: + # all://: + # - http://host.docker.internal:1080 + # + # using_tor_proxy: true + # + # Extra seconds to add in order to account for the time taken by the proxy + # + # extra_proxy_timeout: 10 + # + # uncomment below section only if you have more than one network interface + # which can be the source of outgoing search requests + # + # source_ips: + # - 1.1.1.1 + # - 1.1.1.2 + # - fe80::/126 + +# External plugin configuration, for more details see +# https://docs.searxng.org/dev/plugins.html +# +# plugins: +# - plugin1 +# - plugin2 +# - ... + +# Comment or un-comment plugin to activate / deactivate by default. +# +# enabled_plugins: +# # these plugins are enabled if nothing is configured .. +# - 'Hash plugin' +# - 'Self Information' +# - 'Tracker URL remover' +# - 'Ahmia blacklist' # activation depends on outgoing.using_tor_proxy +# # these plugins are disabled if nothing is configured .. +# - 'Hostnames plugin' # see 'hostnames' configuration below +# - 'Basic Calculator' +# - 'Open Access DOI rewrite' +# - 'Tor check plugin' +# # Read the docs before activate: auto-detection of the language could be +# # detrimental to users expectations / users can activate the plugin in the +# # preferences if they want. +# - 'Autodetect search language' + +# Configuration of the "Hostnames plugin": +# +# hostnames: +# replace: +# '(.*\.)?youtube\.com$': 'invidious.example.com' +# '(.*\.)?youtu\.be$': 'invidious.example.com' +# '(.*\.)?reddit\.com$': 'teddit.example.com' +# '(.*\.)?redd\.it$': 'teddit.example.com' +# '(www\.)?twitter\.com$': 'nitter.example.com' +# remove: +# - '(.*\.)?facebook.com$' +# low_priority: +# - '(.*\.)?google(\..*)?$' +# high_priority: +# - '(.*\.)?wikipedia.org$' +# +# Alternatively you can use external files for configuring the "Hostnames plugin": +# +# hostnames: +# replace: 'rewrite-hosts.yml' +# +# Content of 'rewrite-hosts.yml' (place the file in the same directory as 'settings.yml'): +# '(.*\.)?youtube\.com$': 'invidious.example.com' +# '(.*\.)?youtu\.be$': 'invidious.example.com' +# + +checker: + # disable checker when in debug mode + off_when_debug: true + + # use "scheduling: false" to disable scheduling + # scheduling: interval or int + + # to activate the scheduler: + # * uncomment "scheduling" section + # * add "cache2 = name=searxngcache,items=2000,blocks=2000,blocksize=4096,bitmap=1" + # to your uwsgi.ini + + # scheduling: + # start_after: [300, 1800] # delay to start the first run of the checker + # every: [86400, 90000] # how often the checker runs + + # additional tests: only for the YAML anchors (see the engines section) + # + additional_tests: + rosebud: &test_rosebud + matrix: + query: rosebud + lang: en + result_container: + - not_empty + - ['one_title_contains', 'citizen kane'] + test: + - unique_results + + android: &test_android + matrix: + query: ['android'] + lang: ['en', 'de', 'fr', 'zh-CN'] + result_container: + - not_empty + - ['one_title_contains', 'google'] + test: + - unique_results + + # tests: only for the YAML anchors (see the engines section) + tests: + infobox: &tests_infobox + infobox: + matrix: + query: ["linux", "new york", "bbc"] + result_container: + - has_infobox + +categories_as_tabs: + general: + images: + videos: + news: + map: + music: + it: + science: + files: + social media: + +engines: + - name: 9gag + engine: 9gag + shortcut: 9g + disabled: true + + - name: alpine linux packages + engine: alpinelinux + disabled: true + shortcut: alp + + - name: annas archive + engine: annas_archive + disabled: true + shortcut: aa + + # - name: annas articles + # engine: annas_archive + # shortcut: aaa + # # https://docs.searxng.org/dev/engines/online/annas_archive.html + # aa_content: 'magazine' # book_fiction, book_unknown, book_nonfiction, book_comic + # aa_ext: 'pdf' # pdf, epub, .. + # aa_sort: oldest' # newest, oldest, largest, smallest + + - name: apk mirror + engine: apkmirror + timeout: 4.0 + shortcut: apkm + disabled: true + + - name: apple app store + engine: apple_app_store + shortcut: aps + disabled: true + + # Requires Tor + - name: ahmia + engine: ahmia + categories: onions + enable_http: true + shortcut: ah + + - name: anaconda + engine: xpath + paging: true + first_page_num: 0 + search_url: https://anaconda.org/search?q={query}&page={pageno} + results_xpath: //tbody/tr + url_xpath: ./td/h5/a[last()]/@href + title_xpath: ./td/h5 + content_xpath: ./td[h5]/text() + categories: it + timeout: 6.0 + shortcut: conda + disabled: true + + - name: arch linux wiki + engine: archlinux + shortcut: al + + - name: artic + engine: artic + shortcut: arc + timeout: 4.0 + + - name: arxiv + engine: arxiv + shortcut: arx + timeout: 4.0 + + - name: ask + engine: ask + shortcut: ask + disabled: true + + # tmp suspended: dh key too small + # - name: base + # engine: base + # shortcut: bs + + - name: bandcamp + engine: bandcamp + shortcut: bc + categories: music + + - name: wikipedia + engine: wikipedia + shortcut: wp + # add "list" to the array to get results in the results list + display_type: ["infobox"] + base_url: 'https://{language}.wikipedia.org/' + categories: [general] + + - name: bilibili + engine: bilibili + shortcut: bil + disabled: true + + - name: bing + engine: bing + shortcut: bi + disabled: false + + - name: bing images + engine: bing_images + shortcut: bii + + - name: bing news + engine: bing_news + shortcut: bin + + - name: bing videos + engine: bing_videos + shortcut: biv + + - name: bitbucket + engine: xpath + paging: true + search_url: https://bitbucket.org/repo/all/{pageno}?name={query} + url_xpath: //article[@class="repo-summary"]//a[@class="repo-link"]/@href + title_xpath: //article[@class="repo-summary"]//a[@class="repo-link"] + content_xpath: //article[@class="repo-summary"]/p + categories: [it, repos] + timeout: 4.0 + disabled: true + shortcut: bb + about: + website: https://bitbucket.org/ + wikidata_id: Q2493781 + official_api_documentation: https://developer.atlassian.com/bitbucket + use_official_api: false + require_api_key: false + results: HTML + + - name: bpb + engine: bpb + shortcut: bpb + disabled: true + + - name: btdigg + engine: btdigg + shortcut: bt + disabled: true + + - name: openverse + engine: openverse + categories: images + shortcut: opv + + - name: media.ccc.de + engine: ccc_media + shortcut: c3tv + # We don't set language: de here because media.ccc.de is not just + # for a German audience. It contains many English videos and many + # German videos have English subtitles. + disabled: true + + - name: chefkoch + engine: chefkoch + shortcut: chef + # to show premium or plus results too: + # skip_premium: false + + # - name: core.ac.uk + # engine: core + # categories: science + # shortcut: cor + # # get your API key from: https://core.ac.uk/api-keys/register/ + # api_key: 'unset' + + - name: cppreference + engine: cppreference + shortcut: cpp + paging: false + disabled: true + + - name: crossref + engine: crossref + shortcut: cr + timeout: 30 + disabled: true + + - name: crowdview + engine: json_engine + shortcut: cv + categories: general + paging: false + search_url: https://crowdview-next-js.onrender.com/api/search-v3?query={query} + results_query: results + url_query: link + title_query: title + content_query: snippet + disabled: true + about: + website: https://crowdview.ai/ + + - name: yep + engine: yep + shortcut: yep + categories: general + search_type: web + timeout: 5 + disabled: true + + - name: yep images + engine: yep + shortcut: yepi + categories: images + search_type: images + disabled: true + + - name: yep news + engine: yep + shortcut: yepn + categories: news + search_type: news + disabled: true + + - name: curlie + engine: xpath + shortcut: cl + categories: general + disabled: true + paging: true + lang_all: '' + search_url: https://curlie.org/search?q={query}&lang={lang}&start={pageno}&stime=92452189 + page_size: 20 + results_xpath: //div[@id="site-list-content"]/div[@class="site-item"] + url_xpath: ./div[@class="title-and-desc"]/a/@href + title_xpath: ./div[@class="title-and-desc"]/a/div + content_xpath: ./div[@class="title-and-desc"]/div[@class="site-descr"] + about: + website: https://curlie.org/ + wikidata_id: Q60715723 + use_official_api: false + require_api_key: false + results: HTML + + - name: currency + engine: currency_convert + categories: general + shortcut: cc + + - name: bahnhof + engine: json_engine + search_url: https://www.bahnhof.de/api/stations/search/{query} + url_prefix: https://www.bahnhof.de/ + url_query: slug + title_query: name + content_query: state + shortcut: bf + disabled: true + about: + website: https://www.bahn.de + wikidata_id: Q22811603 + use_official_api: false + require_api_key: false + results: JSON + language: de + tests: + bahnhof: + matrix: + query: berlin + lang: en + result_container: + - not_empty + - ['one_title_contains', 'Berlin Hauptbahnhof'] + test: + - unique_results + + - name: deezer + engine: deezer + shortcut: dz + disabled: true + + - name: destatis + engine: destatis + shortcut: destat + disabled: true + + - name: deviantart + engine: deviantart + shortcut: da + timeout: 3.0 + + - name: ddg definitions + engine: duckduckgo_definitions + shortcut: ddd + weight: 2 + disabled: true + tests: *tests_infobox + + # cloudflare protected + # - name: digbt + # engine: digbt + # shortcut: dbt + # timeout: 6.0 + # disabled: true + + - name: docker hub + engine: docker_hub + shortcut: dh + categories: [it, packages] + + - name: encyclosearch + engine: json_engine + shortcut: es + categories: general + paging: true + search_url: https://encyclosearch.org/encyclosphere/search?q={query}&page={pageno}&resultsPerPage=15 + results_query: Results + url_query: SourceURL + title_query: Title + content_query: Description + disabled: true + about: + website: https://encyclosearch.org + official_api_documentation: https://encyclosearch.org/docs/#/rest-api + use_official_api: true + require_api_key: false + results: JSON + + - name: erowid + engine: xpath + paging: true + first_page_num: 0 + page_size: 30 + search_url: https://www.erowid.org/search.php?q={query}&s={pageno} + url_xpath: //dl[@class="results-list"]/dt[@class="result-title"]/a/@href + title_xpath: //dl[@class="results-list"]/dt[@class="result-title"]/a/text() + content_xpath: //dl[@class="results-list"]/dd[@class="result-details"] + categories: [] + shortcut: ew + disabled: true + about: + website: https://www.erowid.org/ + wikidata_id: Q1430691 + official_api_documentation: + use_official_api: false + require_api_key: false + results: HTML + + # - name: elasticsearch + # shortcut: es + # engine: elasticsearch + # base_url: http://localhost:9200 + # username: elastic + # password: changeme + # index: my-index + # # available options: match, simple_query_string, term, terms, custom + # query_type: match + # # if query_type is set to custom, provide your query here + # #custom_query_json: {"query":{"match_all": {}}} + # #show_metadata: false + # disabled: true + + - name: wikidata + engine: wikidata + shortcut: wd + timeout: 3.0 + weight: 2 + # add "list" to the array to get results in the results list + display_type: ["infobox"] + tests: *tests_infobox + categories: [general] + + - name: duckduckgo + engine: duckduckgo + shortcut: ddg + + - name: duckduckgo images + engine: duckduckgo_extra + categories: [images, web] + ddg_category: images + shortcut: ddi + disabled: true + + - name: duckduckgo videos + engine: duckduckgo_extra + categories: [videos, web] + ddg_category: videos + shortcut: ddv + disabled: true + + - name: duckduckgo news + engine: duckduckgo_extra + categories: [news, web] + ddg_category: news + shortcut: ddn + disabled: true + + - name: duckduckgo weather + engine: duckduckgo_weather + shortcut: ddw + disabled: true + + - name: apple maps + engine: apple_maps + shortcut: apm + disabled: true + timeout: 5.0 + + - name: emojipedia + engine: emojipedia + timeout: 4.0 + shortcut: em + disabled: true + + - name: tineye + engine: tineye + shortcut: tin + timeout: 9.0 + disabled: true + + - name: etymonline + engine: xpath + paging: true + search_url: https://etymonline.com/search?page={pageno}&q={query} + url_xpath: //a[contains(@class, "word__name--")]/@href + title_xpath: //a[contains(@class, "word__name--")] + content_xpath: //section[contains(@class, "word__defination")] + first_page_num: 1 + shortcut: et + categories: [dictionaries] + about: + website: https://www.etymonline.com/ + wikidata_id: Q1188617 + official_api_documentation: + use_official_api: false + require_api_key: false + results: HTML + + # - name: ebay + # engine: ebay + # shortcut: eb + # base_url: 'https://www.ebay.com' + # disabled: true + # timeout: 5 + + - name: 1x + engine: www1x + shortcut: 1x + timeout: 3.0 + disabled: true + + - name: fdroid + engine: fdroid + shortcut: fd + disabled: true + + - name: findthatmeme + engine: findthatmeme + shortcut: ftm + disabled: true + + - name: flickr + categories: images + shortcut: fl + # You can use the engine using the official stable API, but you need an API + # key, see: https://www.flickr.com/services/apps/create/ + # engine: flickr + # api_key: 'apikey' # required! + # Or you can use the html non-stable engine, activated by default + engine: flickr_noapi + + - name: free software directory + engine: mediawiki + shortcut: fsd + categories: [it, software wikis] + base_url: https://directory.fsf.org/ + search_type: title + timeout: 5.0 + disabled: true + about: + website: https://directory.fsf.org/ + wikidata_id: Q2470288 + + # - name: freesound + # engine: freesound + # shortcut: fnd + # disabled: true + # timeout: 15.0 + # API key required, see: https://freesound.org/docs/api/overview.html + # api_key: MyAPIkey + + - name: frinkiac + engine: frinkiac + shortcut: frk + disabled: true + + - name: fyyd + engine: fyyd + shortcut: fy + timeout: 8.0 + disabled: true + + - name: geizhals + engine: geizhals + shortcut: geiz + disabled: true + + - name: genius + engine: genius + shortcut: gen + + - name: gentoo + engine: mediawiki + shortcut: ge + categories: ["it", "software wikis"] + base_url: "https://wiki.gentoo.org/" + api_path: "api.php" + search_type: text + timeout: 10 + + - name: gitlab + engine: json_engine + paging: true + search_url: https://gitlab.com/api/v4/projects?search={query}&page={pageno} + url_query: web_url + title_query: name_with_namespace + content_query: description + page_size: 20 + categories: [it, repos] + shortcut: gl + timeout: 10.0 + disabled: true + about: + website: https://about.gitlab.com/ + wikidata_id: Q16639197 + official_api_documentation: https://docs.gitlab.com/ee/api/ + use_official_api: false + require_api_key: false + results: JSON + + - name: github + engine: github + shortcut: gh + + - name: codeberg + # https://docs.searxng.org/dev/engines/online/gitea.html + engine: gitea + base_url: https://codeberg.org + shortcut: cb + disabled: true + + - name: gitea.com + engine: gitea + base_url: https://gitea.com + shortcut: gitea + disabled: true + + - name: goodreads + engine: goodreads + shortcut: good + timeout: 4.0 + disabled: true + + - name: google + engine: google + shortcut: go + # additional_tests: + # android: *test_android + + - name: google images + engine: google_images + shortcut: goi + # additional_tests: + # android: *test_android + # dali: + # matrix: + # query: ['Dali Christ'] + # lang: ['en', 'de', 'fr', 'zh-CN'] + # result_container: + # - ['one_title_contains', 'Salvador'] + + - name: google news + engine: google_news + shortcut: gon + # additional_tests: + # android: *test_android + + - name: google videos + engine: google_videos + shortcut: gov + # additional_tests: + # android: *test_android + + - name: google scholar + engine: google_scholar + shortcut: gos + + - name: google play apps + engine: google_play + categories: [files, apps] + shortcut: gpa + play_categ: apps + disabled: true + + - name: google play movies + engine: google_play + categories: videos + shortcut: gpm + play_categ: movies + disabled: true + + - name: material icons + engine: material_icons + categories: images + shortcut: mi + disabled: true + + - name: gpodder + engine: json_engine + shortcut: gpod + timeout: 4.0 + paging: false + search_url: https://gpodder.net/search.json?q={query} + url_query: url + title_query: title + content_query: description + page_size: 19 + categories: music + disabled: true + about: + website: https://gpodder.net + wikidata_id: Q3093354 + official_api_documentation: https://gpoddernet.readthedocs.io/en/latest/api/ + use_official_api: false + requires_api_key: false + results: JSON + + - name: habrahabr + engine: xpath + paging: true + search_url: https://habr.com/en/search/page{pageno}/?q={query} + results_xpath: //article[contains(@class, "tm-articles-list__item")] + url_xpath: .//a[@class="tm-title__link"]/@href + title_xpath: .//a[@class="tm-title__link"] + content_xpath: .//div[contains(@class, "article-formatted-body")] + categories: it + timeout: 4.0 + disabled: true + shortcut: habr + about: + website: https://habr.com/ + wikidata_id: Q4494434 + official_api_documentation: https://habr.com/en/docs/help/api/ + use_official_api: false + require_api_key: false + results: HTML + + - name: hackernews + engine: hackernews + shortcut: hn + disabled: true + + - name: hex + engine: hex + shortcut: hex + disabled: true + # Valid values: name inserted_at updated_at total_downloads recent_downloads + sort_criteria: "recent_downloads" + page_size: 10 + + - name: crates.io + engine: crates + shortcut: crates + disabled: true + timeout: 6.0 + + - name: hoogle + engine: xpath + search_url: https://hoogle.haskell.org/?hoogle={query} + results_xpath: '//div[@class="result"]' + title_xpath: './/div[@class="ans"]//a' + url_xpath: './/div[@class="ans"]//a/@href' + content_xpath: './/div[@class="from"]' + page_size: 20 + categories: [it, packages] + shortcut: ho + about: + website: https://hoogle.haskell.org/ + wikidata_id: Q34010 + official_api_documentation: https://hackage.haskell.org/api + use_official_api: false + require_api_key: false + results: JSON + + - name: imdb + engine: imdb + shortcut: imdb + timeout: 6.0 + disabled: true + + - name: imgur + engine: imgur + shortcut: img + disabled: true + + - name: ina + engine: ina + shortcut: in + timeout: 6.0 + disabled: true + + - name: invidious + engine: invidious + # Instanes will be selected randomly, see https://api.invidious.io/ for + # instances that are stable (good uptime) and close to you. + base_url: + - https://invidious.io.lol + - https://invidious.fdn.fr + - https://yt.artemislena.eu + - https://invidious.tiekoetter.com + - https://invidious.flokinet.to + - https://vid.puffyan.us + - https://invidious.privacydev.net + - https://inv.tux.pizza + shortcut: iv + timeout: 3.0 + disabled: true + + - name: jisho + engine: jisho + shortcut: js + timeout: 3.0 + disabled: true + + - name: kickass + engine: kickass + base_url: + - https://kickasstorrents.to + - https://kickasstorrents.cr + - https://kickasstorrent.cr + - https://kickass.sx + - https://kat.am + shortcut: kc + timeout: 4.0 + disabled: true + + - name: lemmy communities + engine: lemmy + lemmy_type: Communities + shortcut: leco + + - name: lemmy users + engine: lemmy + network: lemmy communities + lemmy_type: Users + shortcut: leus + + - name: lemmy posts + engine: lemmy + network: lemmy communities + lemmy_type: Posts + shortcut: lepo + + - name: lemmy comments + engine: lemmy + network: lemmy communities + lemmy_type: Comments + shortcut: lecom + + - name: library genesis + engine: xpath + # search_url: https://libgen.is/search.php?req={query} + search_url: https://libgen.rs/search.php?req={query} + url_xpath: //a[contains(@href,"book/index.php?md5")]/@href + title_xpath: //a[contains(@href,"book/")]/text()[1] + content_xpath: //td/a[1][contains(@href,"=author")]/text() + categories: files + timeout: 7.0 + disabled: true + shortcut: lg + about: + website: https://libgen.fun/ + wikidata_id: Q22017206 + official_api_documentation: + use_official_api: false + require_api_key: false + results: HTML + + - name: z-library + engine: zlibrary + shortcut: zlib + categories: files + timeout: 7.0 + disabled: true + + - name: library of congress + engine: loc + shortcut: loc + categories: images + + - name: libretranslate + engine: libretranslate + # https://github.com/LibreTranslate/LibreTranslate?tab=readme-ov-file#mirrors + base_url: + - https://translate.terraprint.co + - https://trans.zillyhuhn.com + # api_key: abc123 + shortcut: lt + disabled: true + + - name: lingva + engine: lingva + shortcut: lv + # set lingva instance in url, by default it will use the official instance + # url: https://lingva.thedaviddelta.com + + - name: lobste.rs + engine: xpath + search_url: https://lobste.rs/search?q={query}&what=stories&order=relevance + results_xpath: //li[contains(@class, "story")] + url_xpath: .//a[@class="u-url"]/@href + title_xpath: .//a[@class="u-url"] + content_xpath: .//a[@class="domain"] + categories: it + shortcut: lo + timeout: 5.0 + disabled: true + about: + website: https://lobste.rs/ + wikidata_id: Q60762874 + official_api_documentation: + use_official_api: false + require_api_key: false + results: HTML + + - name: mastodon users + engine: mastodon + mastodon_type: accounts + base_url: https://mastodon.social + shortcut: mau + + - name: mastodon hashtags + engine: mastodon + mastodon_type: hashtags + base_url: https://mastodon.social + shortcut: mah + + # - name: matrixrooms + # engine: mrs + # # https://docs.searxng.org/dev/engines/online/mrs.html + # # base_url: https://mrs-api-host + # shortcut: mtrx + # disabled: true + + - name: mdn + shortcut: mdn + engine: json_engine + categories: [it] + paging: true + search_url: https://developer.mozilla.org/api/v1/search?q={query}&page={pageno} + results_query: documents + url_query: mdn_url + url_prefix: https://developer.mozilla.org + title_query: title + content_query: summary + about: + website: https://developer.mozilla.org + wikidata_id: Q3273508 + official_api_documentation: null + use_official_api: false + require_api_key: false + results: JSON + + - name: metacpan + engine: metacpan + shortcut: cpan + disabled: true + number_of_results: 20 + + # - name: meilisearch + # engine: meilisearch + # shortcut: mes + # enable_http: true + # base_url: http://localhost:7700 + # index: my-index + + - name: mixcloud + engine: mixcloud + shortcut: mc + + # MongoDB engine + # Required dependency: pymongo + # - name: mymongo + # engine: mongodb + # shortcut: md + # exact_match_only: false + # host: '127.0.0.1' + # port: 27017 + # enable_http: true + # results_per_page: 20 + # database: 'business' + # collection: 'reviews' # name of the db collection + # key: 'name' # key in the collection to search for + + - name: mozhi + engine: mozhi + base_url: + - https://mozhi.aryak.me + - https://translate.bus-hit.me + - https://nyc1.mz.ggtyler.dev + # mozhi_engine: google - see https://mozhi.aryak.me for supported engines + timeout: 4.0 + shortcut: mz + disabled: true + + - name: mwmbl + engine: mwmbl + # api_url: https://api.mwmbl.org + shortcut: mwm + disabled: true + + - name: npm + engine: npm + shortcut: npm + timeout: 5.0 + disabled: true + + - name: nyaa + engine: nyaa + shortcut: nt + disabled: true + + - name: mankier + engine: json_engine + search_url: https://www.mankier.com/api/v2/mans/?q={query} + results_query: results + url_query: url + title_query: name + content_query: description + categories: it + shortcut: man + about: + website: https://www.mankier.com/ + official_api_documentation: https://www.mankier.com/api + use_official_api: true + require_api_key: false + results: JSON + + # read https://docs.searxng.org/dev/engines/online/mullvad_leta.html + # - name: mullvadleta + # engine: mullvad_leta + # leta_engine: google # choose one of the following: google, brave + # use_cache: true # Only 100 non-cache searches per day, suggested only for private instances + # search_url: https://leta.mullvad.net + # categories: [general, web] + # shortcut: ml + + - name: odysee + engine: odysee + shortcut: od + disabled: true + + - name: openairedatasets + engine: json_engine + paging: true + search_url: https://api.openaire.eu/search/datasets?format=json&page={pageno}&size=10&title={query} + results_query: response/results/result + url_query: metadata/oaf:entity/oaf:result/children/instance/webresource/url/$ + title_query: metadata/oaf:entity/oaf:result/title/$ + content_query: metadata/oaf:entity/oaf:result/description/$ + content_html_to_text: true + categories: "science" + shortcut: oad + timeout: 5.0 + about: + website: https://www.openaire.eu/ + wikidata_id: Q25106053 + official_api_documentation: https://api.openaire.eu/ + use_official_api: false + require_api_key: false + results: JSON + + - name: openairepublications + engine: json_engine + paging: true + search_url: https://api.openaire.eu/search/publications?format=json&page={pageno}&size=10&title={query} + results_query: response/results/result + url_query: metadata/oaf:entity/oaf:result/children/instance/webresource/url/$ + title_query: metadata/oaf:entity/oaf:result/title/$ + content_query: metadata/oaf:entity/oaf:result/description/$ + content_html_to_text: true + categories: science + shortcut: oap + timeout: 5.0 + about: + website: https://www.openaire.eu/ + wikidata_id: Q25106053 + official_api_documentation: https://api.openaire.eu/ + use_official_api: false + require_api_key: false + results: JSON + + - name: openmeteo + engine: open_meteo + shortcut: om + disabled: true + + # - name: opensemanticsearch + # engine: opensemantic + # shortcut: oss + # base_url: 'http://localhost:8983/solr/opensemanticsearch/' + + - name: openstreetmap + engine: openstreetmap + shortcut: osm + + - name: openrepos + engine: xpath + paging: true + search_url: https://openrepos.net/search/node/{query}?page={pageno} + url_xpath: //li[@class="search-result"]//h3[@class="title"]/a/@href + title_xpath: //li[@class="search-result"]//h3[@class="title"]/a + content_xpath: //li[@class="search-result"]//div[@class="search-snippet-info"]//p[@class="search-snippet"] + categories: files + timeout: 4.0 + disabled: true + shortcut: or + about: + website: https://openrepos.net/ + wikidata_id: + official_api_documentation: + use_official_api: false + require_api_key: false + results: HTML + + - name: packagist + engine: json_engine + paging: true + search_url: https://packagist.org/search.json?q={query}&page={pageno} + results_query: results + url_query: url + title_query: name + content_query: description + categories: [it, packages] + disabled: true + timeout: 5.0 + shortcut: pack + about: + website: https://packagist.org + wikidata_id: Q108311377 + official_api_documentation: https://packagist.org/apidoc + use_official_api: true + require_api_key: false + results: JSON + + - name: pdbe + engine: pdbe + shortcut: pdb + # Hide obsolete PDB entries. Default is not to hide obsolete structures + # hide_obsolete: false + + - name: photon + engine: photon + shortcut: ph + + - name: pinterest + engine: pinterest + shortcut: pin + + - name: piped + engine: piped + shortcut: ppd + categories: videos + piped_filter: videos + timeout: 3.0 + + # URL to use as link and for embeds + frontend_url: https://srv.piped.video + # Instance will be selected randomly, for more see https://piped-instances.kavin.rocks/ + backend_url: + - https://pipedapi.kavin.rocks + - https://pipedapi-libre.kavin.rocks + - https://pipedapi.adminforge.de + + - name: piped.music + engine: piped + network: piped + shortcut: ppdm + categories: music + piped_filter: music_songs + timeout: 3.0 + + - name: piratebay + engine: piratebay + shortcut: tpb + # You may need to change this URL to a proxy if piratebay is blocked in your + # country + url: https://thepiratebay.org/ + timeout: 3.0 + + - name: pixiv + shortcut: pv + engine: pixiv + disabled: true + inactive: true + pixiv_image_proxies: + - https://pximg.example.org + # A proxy is required to load the images. Hosting an image proxy server + # for Pixiv: + # --> https://pixivfe.pages.dev/hosting-image-proxy-server/ + # Proxies from public instances. Ask the public instances owners if they + # agree to receive traffic from SearXNG! + # --> https://codeberg.org/VnPower/PixivFE#instances + # --> https://github.com/searxng/searxng/pull/3192#issuecomment-1941095047 + # image proxy of https://pixiv.cat + # - https://i.pixiv.cat + # image proxy of https://www.pixiv.pics + # - https://pximg.cocomi.eu.org + # image proxy of https://pixivfe.exozy.me + # - https://pximg.exozy.me + # image proxy of https://pixivfe.ducks.party + # - https://pixiv.ducks.party + # image proxy of https://pixiv.perennialte.ch + # - https://pximg.perennialte.ch + + - name: podcastindex + engine: podcastindex + shortcut: podcast + + # Required dependency: psychopg2 + # - name: postgresql + # engine: postgresql + # database: postgres + # username: postgres + # password: postgres + # limit: 10 + # query_str: 'SELECT * from my_table WHERE my_column = %(query)s' + # shortcut : psql + + - name: presearch + engine: presearch + search_type: search + categories: [general, web] + shortcut: ps + timeout: 4.0 + disabled: true + + - name: presearch images + engine: presearch + network: presearch + search_type: images + categories: [images, web] + timeout: 4.0 + shortcut: psimg + disabled: true + + - name: presearch videos + engine: presearch + network: presearch + search_type: videos + categories: [general, web] + timeout: 4.0 + shortcut: psvid + disabled: true + + - name: presearch news + engine: presearch + network: presearch + search_type: news + categories: [news, web] + timeout: 4.0 + shortcut: psnews + disabled: true + + - name: pub.dev + engine: xpath + shortcut: pd + search_url: https://pub.dev/packages?q={query}&page={pageno} + paging: true + results_xpath: //div[contains(@class,"packages-item")] + url_xpath: ./div/h3/a/@href + title_xpath: ./div/h3/a + content_xpath: ./div/div/div[contains(@class,"packages-description")]/span + categories: [packages, it] + timeout: 3.0 + disabled: true + first_page_num: 1 + about: + website: https://pub.dev/ + official_api_documentation: https://pub.dev/help/api + use_official_api: false + require_api_key: false + results: HTML + + - name: pubmed + engine: pubmed + shortcut: pub + timeout: 3.0 + + - name: pypi + shortcut: pypi + engine: pypi + + - name: qwant + qwant_categ: web + engine: qwant + disabled: true + shortcut: qw + categories: [general, web] + additional_tests: + rosebud: *test_rosebud + + - name: qwant news + qwant_categ: news + engine: qwant + shortcut: qwn + categories: news + network: qwant + + - name: qwant images + qwant_categ: images + engine: qwant + shortcut: qwi + categories: [images, web] + network: qwant + + - name: qwant videos + qwant_categ: videos + engine: qwant + shortcut: qwv + categories: [videos, web] + network: qwant + + # - name: library + # engine: recoll + # shortcut: lib + # base_url: 'https://recoll.example.org/' + # search_dir: '' + # mount_prefix: /export + # dl_prefix: 'https://download.example.org' + # timeout: 30.0 + # categories: files + # disabled: true + + # - name: recoll library reference + # engine: recoll + # base_url: 'https://recoll.example.org/' + # search_dir: reference + # mount_prefix: /export + # dl_prefix: 'https://download.example.org' + # shortcut: libr + # timeout: 30.0 + # categories: files + # disabled: true + + - name: radio browser + engine: radio_browser + shortcut: rb + + - name: reddit + engine: reddit + shortcut: re + page_size: 25 + disabled: true + + - name: rottentomatoes + engine: rottentomatoes + shortcut: rt + disabled: true + + # Required dependency: redis + # - name: myredis + # shortcut : rds + # engine: redis_server + # exact_match_only: false + # host: '127.0.0.1' + # port: 6379 + # enable_http: true + # password: '' + # db: 0 + + # tmp suspended: bad certificate + # - name: scanr structures + # shortcut: scs + # engine: scanr_structures + # disabled: true + + - name: searchmysite + engine: xpath + shortcut: sms + categories: general + paging: true + search_url: https://searchmysite.net/search/?q={query}&page={pageno} + results_xpath: //div[contains(@class,'search-result')] + url_xpath: .//a[contains(@class,'result-link')]/@href + title_xpath: .//span[contains(@class,'result-title-txt')]/text() + content_xpath: ./p[@id='result-hightlight'] + disabled: true + about: + website: https://searchmysite.net + + - name: sepiasearch + engine: sepiasearch + shortcut: sep + + - name: soundcloud + engine: soundcloud + shortcut: sc + + - name: stackoverflow + engine: stackexchange + shortcut: st + api_site: 'stackoverflow' + categories: [it, q&a] + + - name: askubuntu + engine: stackexchange + shortcut: ubuntu + api_site: 'askubuntu' + categories: [it, q&a] + + - name: internetarchivescholar + engine: internet_archive_scholar + shortcut: ias + timeout: 15.0 + + - name: superuser + engine: stackexchange + shortcut: su + api_site: 'superuser' + categories: [it, q&a] + + - name: discuss.python + engine: discourse + shortcut: dpy + base_url: 'https://discuss.python.org' + categories: [it, q&a] + disabled: true + + - name: caddy.community + engine: discourse + shortcut: caddy + base_url: 'https://caddy.community' + categories: [it, q&a] + disabled: true + + - name: pi-hole.community + engine: discourse + shortcut: pi + categories: [it, q&a] + base_url: 'https://discourse.pi-hole.net' + disabled: true + + - name: searchcode code + engine: searchcode_code + shortcut: scc + disabled: true + + # - name: searx + # engine: searx_engine + # shortcut: se + # instance_urls : + # - http://127.0.0.1:8888/ + # - ... + # disabled: true + + - name: semantic scholar + engine: semantic_scholar + disabled: true + shortcut: se + + # Spotify needs API credentials + # - name: spotify + # engine: spotify + # shortcut: stf + # api_client_id: ******* + # api_client_secret: ******* + + # - name: solr + # engine: solr + # shortcut: slr + # base_url: http://localhost:8983 + # collection: collection_name + # sort: '' # sorting: asc or desc + # field_list: '' # comma separated list of field names to display on the UI + # default_fields: '' # default field to query + # query_fields: '' # query fields + # enable_http: true + + # - name: springer nature + # engine: springer + # # get your API key from: https://dev.springernature.com/signup + # # working API key, for test & debug: "a69685087d07eca9f13db62f65b8f601" + # api_key: 'unset' + # shortcut: springer + # timeout: 15.0 + + - name: startpage + engine: startpage + shortcut: sp + timeout: 6.0 + disabled: true + additional_tests: + rosebud: *test_rosebud + + - name: tokyotoshokan + engine: tokyotoshokan + shortcut: tt + timeout: 6.0 + disabled: true + + - name: solidtorrents + engine: solidtorrents + shortcut: solid + timeout: 4.0 + base_url: + - https://solidtorrents.to + - https://bitsearch.to + + # For this demo of the sqlite engine download: + # https://liste.mediathekview.de/filmliste-v2.db.bz2 + # and unpack into searx/data/filmliste-v2.db + # Query to test: "!demo concert" + # + # - name: demo + # engine: sqlite + # shortcut: demo + # categories: general + # result_template: default.html + # database: searx/data/filmliste-v2.db + # query_str: >- + # SELECT title || ' (' || time(duration, 'unixepoch') || ')' AS title, + # COALESCE( NULLIF(url_video_hd,''), NULLIF(url_video_sd,''), url_video) AS url, + # description AS content + # FROM film + # WHERE title LIKE :wildcard OR description LIKE :wildcard + # ORDER BY duration DESC + + - name: tagesschau + engine: tagesschau + # when set to false, display URLs from Tagesschau, and not the actual source + # (e.g. NDR, WDR, SWR, HR, ...) + use_source_url: true + shortcut: ts + disabled: true + + - name: tmdb + engine: xpath + paging: true + categories: movies + search_url: https://www.themoviedb.org/search?page={pageno}&query={query} + results_xpath: //div[contains(@class,"movie") or contains(@class,"tv")]//div[contains(@class,"card")] + url_xpath: .//div[contains(@class,"poster")]/a/@href + thumbnail_xpath: .//img/@src + title_xpath: .//div[contains(@class,"title")]//h2 + content_xpath: .//div[contains(@class,"overview")] + shortcut: tm + disabled: true + + # Requires Tor + - name: torch + engine: xpath + paging: true + search_url: + http://xmh57jrknzkhv6y3ls3ubitzfqnkrwxhopf5aygthi7d6rplyvk3noyd.onion/cgi-bin/omega/omega?P={query}&DEFAULTOP=and + results_xpath: //table//tr + url_xpath: ./td[2]/a + title_xpath: ./td[2]/b + content_xpath: ./td[2]/small + categories: onions + enable_http: true + shortcut: tch + + # torznab engine lets you query any torznab compatible indexer. Using this + # engine in combination with Jackett opens the possibility to query a lot of + # public and private indexers directly from SearXNG. More details at: + # https://docs.searxng.org/dev/engines/online/torznab.html + # + # - name: Torznab EZTV + # engine: torznab + # shortcut: eztv + # base_url: http://localhost:9117/api/v2.0/indexers/eztv/results/torznab + # enable_http: true # if using localhost + # api_key: xxxxxxxxxxxxxxx + # show_magnet_links: true + # show_torrent_files: false + # # https://github.com/Jackett/Jackett/wiki/Jackett-Categories + # torznab_categories: # optional + # - 2000 + # - 5000 + + # tmp suspended - too slow, too many errors + # - name: urbandictionary + # engine : xpath + # search_url : https://www.urbandictionary.com/define.php?term={query} + # url_xpath : //*[@class="word"]/@href + # title_xpath : //*[@class="def-header"] + # content_xpath: //*[@class="meaning"] + # shortcut: ud + + - name: unsplash + engine: unsplash + shortcut: us + + - name: yandex music + engine: yandex_music + shortcut: ydm + disabled: true + # https://yandex.com/support/music/access.html + inactive: true + + - name: yahoo + engine: yahoo + shortcut: yh + disabled: true + + - name: yahoo news + engine: yahoo_news + shortcut: yhn + + - name: youtube + shortcut: yt + # You can use the engine using the official stable API, but you need an API + # key See: https://console.developers.google.com/project + # + # engine: youtube_api + # api_key: 'apikey' # required! + # + # Or you can use the html non-stable engine, activated by default + engine: youtube_noapi + + - name: dailymotion + engine: dailymotion + shortcut: dm + + - name: vimeo + engine: vimeo + shortcut: vm + disabled: true + + - name: wiby + engine: json_engine + paging: true + search_url: https://wiby.me/json/?q={query}&p={pageno} + url_query: URL + title_query: Title + content_query: Snippet + categories: [general, web] + shortcut: wib + disabled: true + about: + website: https://wiby.me/ + + - name: alexandria + engine: json_engine + shortcut: alx + categories: general + paging: true + search_url: https://api.alexandria.org/?a=1&q={query}&p={pageno} + results_query: results + title_query: title + url_query: url + content_query: snippet + timeout: 1.5 + disabled: true + about: + website: https://alexandria.org/ + official_api_documentation: https://github.com/alexandria-org/alexandria-api/raw/master/README.md + use_official_api: true + require_api_key: false + results: JSON + + - name: wikibooks + engine: mediawiki + weight: 0.5 + shortcut: wb + categories: [general, wikimedia] + base_url: "https://{language}.wikibooks.org/" + search_type: text + disabled: true + about: + website: https://www.wikibooks.org/ + wikidata_id: Q367 + + - name: wikinews + engine: mediawiki + shortcut: wn + categories: [news, wikimedia] + base_url: "https://{language}.wikinews.org/" + search_type: text + srsort: create_timestamp_desc + about: + website: https://www.wikinews.org/ + wikidata_id: Q964 + + - name: wikiquote + engine: mediawiki + weight: 0.5 + shortcut: wq + categories: [general, wikimedia] + base_url: "https://{language}.wikiquote.org/" + search_type: text + disabled: true + additional_tests: + rosebud: *test_rosebud + about: + website: https://www.wikiquote.org/ + wikidata_id: Q369 + + - name: wikisource + engine: mediawiki + weight: 0.5 + shortcut: ws + categories: [general, wikimedia] + base_url: "https://{language}.wikisource.org/" + search_type: text + disabled: true + about: + website: https://www.wikisource.org/ + wikidata_id: Q263 + + - name: wikispecies + engine: mediawiki + shortcut: wsp + categories: [general, science, wikimedia] + base_url: "https://species.wikimedia.org/" + search_type: text + disabled: true + about: + website: https://species.wikimedia.org/ + wikidata_id: Q13679 + tests: + wikispecies: + matrix: + query: "Campbell, L.I. et al. 2011: MicroRNAs" + lang: en + result_container: + - not_empty + - ['one_title_contains', 'Tardigrada'] + test: + - unique_results + + - name: wiktionary + engine: mediawiki + shortcut: wt + categories: [dictionaries, wikimedia] + base_url: "https://{language}.wiktionary.org/" + search_type: text + about: + website: https://www.wiktionary.org/ + wikidata_id: Q151 + + - name: wikiversity + engine: mediawiki + weight: 0.5 + shortcut: wv + categories: [general, wikimedia] + base_url: "https://{language}.wikiversity.org/" + search_type: text + disabled: true + about: + website: https://www.wikiversity.org/ + wikidata_id: Q370 + + - name: wikivoyage + engine: mediawiki + weight: 0.5 + shortcut: wy + categories: [general, wikimedia] + base_url: "https://{language}.wikivoyage.org/" + search_type: text + disabled: true + about: + website: https://www.wikivoyage.org/ + wikidata_id: Q373 + + - name: wikicommons.images + engine: wikicommons + shortcut: wc + categories: images + search_type: images + number_of_results: 10 + + - name: wikicommons.videos + engine: wikicommons + shortcut: wcv + categories: videos + search_type: videos + number_of_results: 10 + + - name: wikicommons.audio + engine: wikicommons + shortcut: wca + categories: music + search_type: audio + number_of_results: 10 + + - name: wikicommons.files + engine: wikicommons + shortcut: wcf + categories: files + search_type: files + number_of_results: 10 + + - name: wolframalpha + shortcut: wa + # You can use the engine using the official stable API, but you need an API + # key. See: https://products.wolframalpha.com/api/ + # + # engine: wolframalpha_api + # api_key: '' + # + # Or you can use the html non-stable engine, activated by default + engine: wolframalpha_noapi + timeout: 6.0 + categories: general + disabled: true + + - name: dictzone + engine: dictzone + shortcut: dc + + - name: mymemory translated + engine: translated + shortcut: tl + timeout: 5.0 + # You can use without an API key, but you are limited to 1000 words/day + # See: https://mymemory.translated.net/doc/usagelimits.php + # api_key: '' + + # Required dependency: mysql-connector-python + # - name: mysql + # engine: mysql_server + # database: mydatabase + # username: user + # password: pass + # limit: 10 + # query_str: 'SELECT * from mytable WHERE fieldname=%(query)s' + # shortcut: mysql + + - name: 1337x + engine: 1337x + shortcut: 1337x + disabled: true + + - name: duden + engine: duden + shortcut: du + disabled: true + + - name: seznam + shortcut: szn + engine: seznam + disabled: true + + # - name: deepl + # engine: deepl + # shortcut: dpl + # # You can use the engine using the official stable API, but you need an API key + # # See: https://www.deepl.com/pro-api?cta=header-pro-api + # api_key: '' # required! + # timeout: 5.0 + # disabled: true + + - name: mojeek + shortcut: mjk + engine: mojeek + categories: [general, web] + disabled: true + + - name: mojeek images + shortcut: mjkimg + engine: mojeek + categories: [images, web] + search_type: images + paging: false + disabled: true + + - name: mojeek news + shortcut: mjknews + engine: mojeek + categories: [news, web] + search_type: news + paging: false + disabled: true + + - name: moviepilot + engine: moviepilot + shortcut: mp + disabled: true + + - name: naver + shortcut: nvr + categories: [general, web] + engine: xpath + paging: true + search_url: https://search.naver.com/search.naver?where=webkr&sm=osp_hty&ie=UTF-8&query={query}&start={pageno} + url_xpath: //a[@class="link_tit"]/@href + title_xpath: //a[@class="link_tit"] + content_xpath: //div[@class="total_dsc_wrap"]/a + first_page_num: 1 + page_size: 10 + disabled: true + about: + website: https://www.naver.com/ + wikidata_id: Q485639 + official_api_documentation: https://developers.naver.com/docs/nmt/examples/ + use_official_api: false + require_api_key: false + results: HTML + language: ko + + - name: rubygems + shortcut: rbg + engine: xpath + paging: true + search_url: https://rubygems.org/search?page={pageno}&query={query} + results_xpath: /html/body/main/div/a[@class="gems__gem"] + url_xpath: ./@href + title_xpath: ./span/h2 + content_xpath: ./span/p + suggestion_xpath: /html/body/main/div/div[@class="search__suggestions"]/p/a + first_page_num: 1 + categories: [it, packages] + disabled: true + about: + website: https://rubygems.org/ + wikidata_id: Q1853420 + official_api_documentation: https://guides.rubygems.org/rubygems-org-api/ + use_official_api: false + require_api_key: false + results: HTML + + - name: peertube + engine: peertube + shortcut: ptb + paging: true + # alternatives see: https://instances.joinpeertube.org/instances + # base_url: https://tube.4aem.com + categories: videos + disabled: true + timeout: 6.0 + + - name: mediathekviewweb + engine: mediathekviewweb + shortcut: mvw + disabled: true + + - name: yacy + # https://docs.searxng.org/dev/engines/online/yacy.html + engine: yacy + categories: general + search_type: text + base_url: + - https://yacy.searchlab.eu + # see https://github.com/searxng/searxng/pull/3631#issuecomment-2240903027 + # - https://search.kyun.li + # - https://yacy.securecomcorp.eu + # - https://yacy.myserv.ca + # - https://yacy.nsupdate.info + # - https://yacy.electroncash.de + shortcut: ya + disabled: true + # if you aren't using HTTPS for your local yacy instance disable https + # enable_http: false + search_mode: 'global' + # timeout can be reduced in 'local' search mode + timeout: 5.0 + + - name: yacy images + engine: yacy + network: yacy + categories: images + search_type: image + shortcut: yai + disabled: true + # timeout can be reduced in 'local' search mode + timeout: 5.0 + + - name: rumble + engine: rumble + shortcut: ru + base_url: https://rumble.com/ + paging: true + categories: videos + disabled: true + + - name: livespace + engine: livespace + shortcut: ls + categories: videos + disabled: true + timeout: 5.0 + + - name: wordnik + engine: wordnik + shortcut: def + base_url: https://www.wordnik.com/ + categories: [dictionaries] + timeout: 5.0 + + - name: woxikon.de synonyme + engine: xpath + shortcut: woxi + categories: [dictionaries] + timeout: 5.0 + disabled: true + search_url: https://synonyme.woxikon.de/synonyme/{query}.php + url_xpath: //div[@class="upper-synonyms"]/a/@href + content_xpath: //div[@class="synonyms-list-group"] + title_xpath: //div[@class="upper-synonyms"]/a + no_result_for_http_status: [404] + about: + website: https://www.woxikon.de/ + wikidata_id: # No Wikidata ID + use_official_api: false + require_api_key: false + results: HTML + language: de + + - name: seekr news + engine: seekr + shortcut: senews + categories: news + seekr_category: news + disabled: true + + - name: seekr images + engine: seekr + network: seekr news + shortcut: seimg + categories: images + seekr_category: images + disabled: true + + - name: seekr videos + engine: seekr + network: seekr news + shortcut: sevid + categories: videos + seekr_category: videos + disabled: true + + - name: sjp.pwn + engine: sjp + shortcut: sjp + base_url: https://sjp.pwn.pl/ + timeout: 5.0 + disabled: true + + - name: stract + engine: stract + shortcut: str + disabled: true + + - name: svgrepo + engine: svgrepo + shortcut: svg + timeout: 10.0 + disabled: true + + - name: tootfinder + engine: tootfinder + shortcut: toot + + - name: voidlinux + engine: voidlinux + shortcut: void + disabled: true + + - name: wallhaven + engine: wallhaven + # api_key: abcdefghijklmnopqrstuvwxyz + shortcut: wh + + # wikimini: online encyclopedia for children + # The fulltext and title parameter is necessary for Wikimini because + # sometimes it will not show the results and redirect instead + - name: wikimini + engine: xpath + shortcut: wkmn + search_url: https://fr.wikimini.org/w/index.php?search={query}&title=Sp%C3%A9cial%3ASearch&fulltext=Search + url_xpath: //li/div[@class="mw-search-result-heading"]/a/@href + title_xpath: //li//div[@class="mw-search-result-heading"]/a + content_xpath: //li/div[@class="searchresult"] + categories: general + disabled: true + about: + website: https://wikimini.org/ + wikidata_id: Q3568032 + use_official_api: false + require_api_key: false + results: HTML + language: fr + + - name: wttr.in + engine: wttr + shortcut: wttr + timeout: 9.0 + + - name: yummly + engine: yummly + shortcut: yum + disabled: true + + - name: brave + engine: brave + shortcut: br + time_range_support: true + paging: true + categories: [general, web] + brave_category: search + # brave_spellcheck: true + + - name: brave.images + engine: brave + network: brave + shortcut: brimg + categories: [images, web] + brave_category: images + + - name: brave.videos + engine: brave + network: brave + shortcut: brvid + categories: [videos, web] + brave_category: videos + + - name: brave.news + engine: brave + network: brave + shortcut: brnews + categories: news + brave_category: news + + # - name: brave.goggles + # engine: brave + # network: brave + # shortcut: brgog + # time_range_support: true + # paging: true + # categories: [general, web] + # brave_category: goggles + # Goggles: # required! This should be a URL ending in .goggle + + - name: lib.rs + shortcut: lrs + engine: lib_rs + disabled: true + + - name: sourcehut + shortcut: srht + engine: xpath + paging: true + search_url: https://sr.ht/projects?page={pageno}&search={query} + results_xpath: (//div[@class="event-list"])[1]/div[@class="event"] + url_xpath: ./h4/a[2]/@href + title_xpath: ./h4/a[2] + content_xpath: ./p + first_page_num: 1 + categories: [it, repos] + disabled: true + about: + website: https://sr.ht + wikidata_id: Q78514485 + official_api_documentation: https://man.sr.ht/ + use_official_api: false + require_api_key: false + results: HTML + + - name: goo + shortcut: goo + engine: xpath + paging: true + search_url: https://search.goo.ne.jp/web.jsp?MT={query}&FR={pageno}0 + url_xpath: //div[@class="result"]/p[@class='title fsL1']/a/@href + title_xpath: //div[@class="result"]/p[@class='title fsL1']/a + content_xpath: //p[contains(@class,'url fsM')]/following-sibling::p + first_page_num: 0 + categories: [general, web] + disabled: true + timeout: 4.0 + about: + website: https://search.goo.ne.jp + wikidata_id: Q249044 + use_official_api: false + require_api_key: false + results: HTML + language: ja + + - name: bt4g + engine: bt4g + shortcut: bt4g + + - name: pkg.go.dev + engine: pkg_go_dev + shortcut: pgo + disabled: true + +# Doku engine lets you access to any Doku wiki instance: +# A public one or a privete/corporate one. +# - name: ubuntuwiki +# engine: doku +# shortcut: uw +# base_url: 'https://doc.ubuntu-fr.org' + +# Be careful when enabling this engine if you are +# running a public instance. Do not expose any sensitive +# information. You can restrict access by configuring a list +# of access tokens under tokens. +# - name: git grep +# engine: command +# command: ['git', 'grep', '{{QUERY}}'] +# shortcut: gg +# tokens: [] +# disabled: true +# delimiter: +# chars: ':' +# keys: ['filepath', 'code'] + +# Be careful when enabling this engine if you are +# running a public instance. Do not expose any sensitive +# information. You can restrict access by configuring a list +# of access tokens under tokens. +# - name: locate +# engine: command +# command: ['locate', '{{QUERY}}'] +# shortcut: loc +# tokens: [] +# disabled: true +# delimiter: +# chars: ' ' +# keys: ['line'] + +# Be careful when enabling this engine if you are +# running a public instance. Do not expose any sensitive +# information. You can restrict access by configuring a list +# of access tokens under tokens. +# - name: find +# engine: command +# command: ['find', '.', '-name', '{{QUERY}}'] +# query_type: path +# shortcut: fnd +# tokens: [] +# disabled: true +# delimiter: +# chars: ' ' +# keys: ['line'] + +# Be careful when enabling this engine if you are +# running a public instance. Do not expose any sensitive +# information. You can restrict access by configuring a list +# of access tokens under tokens. +# - name: pattern search in files +# engine: command +# command: ['fgrep', '{{QUERY}}'] +# shortcut: fgr +# tokens: [] +# disabled: true +# delimiter: +# chars: ' ' +# keys: ['line'] + +# Be careful when enabling this engine if you are +# running a public instance. Do not expose any sensitive +# information. You can restrict access by configuring a list +# of access tokens under tokens. +# - name: regex search in files +# engine: command +# command: ['grep', '{{QUERY}}'] +# shortcut: gr +# tokens: [] +# disabled: true +# delimiter: +# chars: ' ' +# keys: ['line'] + +doi_resolvers: + oadoi.org: 'https://oadoi.org/' + doi.org: 'https://doi.org/' + doai.io: 'https://dissem.in/' + sci-hub.se: 'https://sci-hub.se/' + sci-hub.st: 'https://sci-hub.st/' + sci-hub.ru: 'https://sci-hub.ru/' + +default_doi_resolver: 'oadoi.org' diff --git a/api/core/tools/provider/builtin/searxng/docker/uwsgi.ini b/api/core/tools/provider/builtin/searxng/docker/uwsgi.ini new file mode 100644 index 0000000000000000000000000000000000000000..9db3d762649fc5ebaeb3d914e82d8e6eb5ddaf2d --- /dev/null +++ b/api/core/tools/provider/builtin/searxng/docker/uwsgi.ini @@ -0,0 +1,54 @@ +[uwsgi] +# Who will run the code +uid = searxng +gid = searxng + +# Number of workers (usually CPU count) +# default value: %k (= number of CPU core, see Dockerfile) +workers = %k + +# Number of threads per worker +# default value: 4 (see Dockerfile) +threads = 4 + +# The right granted on the created socket +chmod-socket = 666 + +# Plugin to use and interpreter config +single-interpreter = true +master = true +plugin = python3 +lazy-apps = true +enable-threads = 4 + +# Module to import +module = searx.webapp + +# Virtualenv and python path +pythonpath = /usr/local/searxng/ +chdir = /usr/local/searxng/searx/ + +# automatically set processes name to something meaningful +auto-procname = true + +# Disable request logging for privacy +disable-logging = true +log-5xx = true + +# Set the max size of a request (request-body excluded) +buffer-size = 8192 + +# No keep alive +# See https://github.com/searx/searx-docker/issues/24 +add-header = Connection: close + +# Follow SIGTERM convention +# See https://github.com/searxng/searxng/issues/3427 +die-on-term + +# uwsgi serves the static files +static-map = /static=/usr/local/searxng/searx/static +# expires set to one day +static-expires = /* 86400 +static-gzip-all = True +offload-threads = 4 diff --git a/api/core/tools/provider/builtin/searxng/searxng.py b/api/core/tools/provider/builtin/searxng/searxng.py new file mode 100644 index 0000000000000000000000000000000000000000..b7bbcc60b1ed26cbe4c278f35338c2d52bae39f4 --- /dev/null +++ b/api/core/tools/provider/builtin/searxng/searxng.py @@ -0,0 +1,20 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.searxng.tools.searxng_search import SearXNGSearchTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class SearXNGProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + SearXNGSearchTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={"query": "SearXNG", "limit": 1, "search_type": "general"}, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/searxng/searxng.yaml b/api/core/tools/provider/builtin/searxng/searxng.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9554c93d5a0c53d302190a66727ca097154f308d --- /dev/null +++ b/api/core/tools/provider/builtin/searxng/searxng.yaml @@ -0,0 +1,24 @@ +identity: + author: Junytang + name: searxng + label: + en_US: SearXNG + zh_Hans: SearXNG + description: + en_US: A free internet metasearch engine. + zh_Hans: 开源免费的互联网元搜索引擎 + icon: icon.svg + tags: + - search + - productivity +credentials_for_provider: + searxng_base_url: + type: text-input + required: true + label: + en_US: SearXNG base URL + zh_Hans: SearXNG base URL + placeholder: + en_US: Please input your SearXNG base URL + zh_Hans: 请输入您的 SearXNG base URL + url: https://docs.dify.ai/tutorials/tool-configuration/searxng diff --git a/api/core/tools/provider/builtin/searxng/tools/searxng_search.py b/api/core/tools/provider/builtin/searxng/tools/searxng_search.py new file mode 100644 index 0000000000000000000000000000000000000000..c5e339a108e5b2547fd58d9232fabecef6815198 --- /dev/null +++ b/api/core/tools/provider/builtin/searxng/tools/searxng_search.py @@ -0,0 +1,46 @@ +from typing import Any + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class SearXNGSearchTool(BuiltinTool): + """ + Tool for performing a search using SearXNG engine. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage | list[ToolInvokeMessage]: + """ + Invoke the SearXNG search tool. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (dict[str, Any]): The parameters for the tool invocation. + + Returns: + ToolInvokeMessage | list[ToolInvokeMessage]: The result of the tool invocation. + """ + + host = self.runtime.credentials.get("searxng_base_url") + if not host: + raise Exception("SearXNG api is required") + + response = requests.get( + host, + params={ + "q": tool_parameters.get("query"), + "format": "json", + "categories": tool_parameters.get("search_type", "general"), + }, + ) + + if response.status_code != 200: + raise Exception(f"Error {response.status_code}: {response.text}") + + res = response.json().get("results", []) + if not res: + return self.create_text_message(f"No results found, get response: {response.content}") + + return [self.create_json_message(item) for item in res] diff --git a/api/core/tools/provider/builtin/searxng/tools/searxng_search.yaml b/api/core/tools/provider/builtin/searxng/tools/searxng_search.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a5e448a30375b44c058920d2c2b1e46d66127981 --- /dev/null +++ b/api/core/tools/provider/builtin/searxng/tools/searxng_search.yaml @@ -0,0 +1,69 @@ +identity: + name: searxng_search + author: Junytang + label: + en_US: SearXNG Search + zh_Hans: SearXNG 搜索 +description: + human: + en_US: SearXNG is a free internet metasearch engine which aggregates results from more than 70 search services. + zh_Hans: SearXNG 是一个免费的互联网元搜索引擎,它从70多个不同的搜索服务中聚合搜索结果。 + llm: Perform searches on SearXNG and get results. +parameters: + - name: query + type: string + required: true + label: + en_US: Query string + zh_Hans: 查询语句 + llm_description: Key words for searching + form: llm + - name: search_type + type: select + required: true + label: + en_US: search type + zh_Hans: 搜索类型 + default: general + options: + - value: general + label: + en_US: General + zh_Hans: 综合 + - value: images + label: + en_US: Images + zh_Hans: 图片 + - value: videos + label: + en_US: Videos + zh_Hans: 视频 + - value: news + label: + en_US: News + zh_Hans: 新闻 + - value: map + label: + en_US: Map + zh_Hans: 地图 + - value: music + label: + en_US: Music + zh_Hans: 音乐 + - value: it + label: + en_US: It + zh_Hans: 信息技术 + - value: science + label: + en_US: Science + zh_Hans: 科学 + - value: files + label: + en_US: Files + zh_Hans: 文件 + - value: social_media + label: + en_US: Social Media + zh_Hans: 社交媒体 + form: form diff --git a/api/core/tools/provider/builtin/serper/_assets/icon.svg b/api/core/tools/provider/builtin/serper/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..3f973a552e5e1709f605282b21e68c8582d28c12 --- /dev/null +++ b/api/core/tools/provider/builtin/serper/_assets/icon.svg @@ -0,0 +1,12 @@ + + + serper + + + + + + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/serper/serper.py b/api/core/tools/provider/builtin/serper/serper.py new file mode 100644 index 0000000000000000000000000000000000000000..cb1d090a9dd4b019be5fba9e4cd772feb5dec1c1 --- /dev/null +++ b/api/core/tools/provider/builtin/serper/serper.py @@ -0,0 +1,20 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.serper.tools.serper_search import SerperSearchTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class SerperProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + SerperSearchTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={"query": "test", "result_type": "link"}, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/serper/serper.yaml b/api/core/tools/provider/builtin/serper/serper.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b3b2d76c4b65731fe66fcaa974d44cb6f1a63aab --- /dev/null +++ b/api/core/tools/provider/builtin/serper/serper.yaml @@ -0,0 +1,31 @@ +identity: + author: zhuhao + name: serper + label: + en_US: Serper + zh_Hans: Serper + pt_BR: Serper + description: + en_US: Serper is a powerful real-time search engine tool API that provides structured data from Google Search. + zh_Hans: Serper 是一个强大的实时搜索引擎工具API,可提供来自 Google 搜索引擎搜索的结构化数据。 + pt_BR: Serper is a powerful real-time search engine tool API that provides structured data from Google Search. + icon: icon.svg + tags: + - search +credentials_for_provider: + serperapi_api_key: + type: secret-input + required: true + label: + en_US: Serper API key + zh_Hans: Serper API key + pt_BR: Serper API key + placeholder: + en_US: Please input your Serper API key + zh_Hans: 请输入你的 Serper API key + pt_BR: Please input your Serper API key + help: + en_US: Get your Serper API key from Serper + zh_Hans: 从 Serper 获取您的 Serper API key + pt_BR: Get your Serper API key from Serper + url: https://serper.dev/api-key diff --git a/api/core/tools/provider/builtin/serper/tools/serper_search.py b/api/core/tools/provider/builtin/serper/tools/serper_search.py new file mode 100644 index 0000000000000000000000000000000000000000..7baebbf95855e0545e98b5c7c9da3a672cc58c2a --- /dev/null +++ b/api/core/tools/provider/builtin/serper/tools/serper_search.py @@ -0,0 +1,34 @@ +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +SERPER_API_URL = "https://google.serper.dev/search" + + +class SerperSearchTool(BuiltinTool): + def _parse_response(self, response: dict) -> dict: + result = {} + if "knowledgeGraph" in response: + result["title"] = response["knowledgeGraph"].get("title", "") + result["description"] = response["knowledgeGraph"].get("description", "") + if "organic" in response: + result["organic"] = [ + {"title": item.get("title", ""), "link": item.get("link", ""), "snippet": item.get("snippet", "")} + for item in response["organic"] + ] + return result + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + params = {"q": tool_parameters["query"], "gl": "us", "hl": "en"} + headers = {"X-API-KEY": self.runtime.credentials["serperapi_api_key"], "Content-Type": "application/json"} + response = requests.get(url=SERPER_API_URL, params=params, headers=headers) + response.raise_for_status() + valuable_res = self._parse_response(response.json()) + return self.create_json_message(valuable_res) diff --git a/api/core/tools/provider/builtin/serper/tools/serper_search.yaml b/api/core/tools/provider/builtin/serper/tools/serper_search.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e1c0a056e65513bda845ebd9b9dd632173d57768 --- /dev/null +++ b/api/core/tools/provider/builtin/serper/tools/serper_search.yaml @@ -0,0 +1,27 @@ +identity: + name: serper + author: zhuhao + label: + en_US: Serper + zh_Hans: Serper + pt_BR: Serper +description: + human: + en_US: A tool for performing a Google search and extracting snippets and webpages.Input should be a search query. + zh_Hans: 一个用于执行 Google 搜索并提取片段和网页的工具。输入应该是一个搜索查询。 + pt_BR: A tool for performing a Google search and extracting snippets and webpages.Input should be a search query. + llm: A tool for performing a Google search and extracting snippets and webpages.Input should be a search query. +parameters: + - name: query + type: string + required: true + label: + en_US: Query string + zh_Hans: 查询语句 + pt_BR: Query string + human_description: + en_US: used for searching + zh_Hans: 用于搜索网页内容 + pt_BR: used for searching + llm_description: key words for searching + form: llm diff --git a/api/core/tools/provider/builtin/siliconflow/_assets/icon.svg b/api/core/tools/provider/builtin/siliconflow/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..ad6b384f7acd212ef5d5b9964c4c5cc47ea07367 --- /dev/null +++ b/api/core/tools/provider/builtin/siliconflow/_assets/icon.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/siliconflow/siliconflow.py b/api/core/tools/provider/builtin/siliconflow/siliconflow.py new file mode 100644 index 0000000000000000000000000000000000000000..37a0b0755b1d39c9f5198a77f5419963820cf30d --- /dev/null +++ b/api/core/tools/provider/builtin/siliconflow/siliconflow.py @@ -0,0 +1,17 @@ +import requests + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class SiliconflowProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + url = "https://api.siliconflow.cn/v1/models" + headers = { + "accept": "application/json", + "authorization": f"Bearer {credentials.get('siliconFlow_api_key')}", + } + + response = requests.get(url, headers=headers) + if response.status_code != 200: + raise ToolProviderCredentialValidationError("SiliconFlow API key is invalid") diff --git a/api/core/tools/provider/builtin/siliconflow/siliconflow.yaml b/api/core/tools/provider/builtin/siliconflow/siliconflow.yaml new file mode 100644 index 0000000000000000000000000000000000000000..46be99f262f2116e2b9a51f97bb327ef168ac8c6 --- /dev/null +++ b/api/core/tools/provider/builtin/siliconflow/siliconflow.yaml @@ -0,0 +1,21 @@ +identity: + author: hjlarry + name: siliconflow + label: + en_US: SiliconFlow + zh_CN: 硅基流动 + description: + en_US: The image generation API provided by SiliconFlow includes Flux and Stable Diffusion models. + zh_CN: 硅基流动提供的图片生成 API,包含 Flux 和 Stable Diffusion 模型。 + icon: icon.svg + tags: + - image +credentials_for_provider: + siliconFlow_api_key: + type: secret-input + required: true + label: + en_US: SiliconFlow API Key + placeholder: + en_US: Please input your SiliconFlow API key + url: https://cloud.siliconflow.cn/account/ak diff --git a/api/core/tools/provider/builtin/siliconflow/tools/flux.py b/api/core/tools/provider/builtin/siliconflow/tools/flux.py new file mode 100644 index 0000000000000000000000000000000000000000..0d16ff385eb30d98ab96dff9c981888ea02fe51d --- /dev/null +++ b/api/core/tools/provider/builtin/siliconflow/tools/flux.py @@ -0,0 +1,43 @@ +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +FLUX_URL = { + "schnell": "https://api.siliconflow.cn/v1/black-forest-labs/FLUX.1-schnell/text-to-image", + "dev": "https://api.siliconflow.cn/v1/image/generations", +} + + +class FluxTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + headers = { + "accept": "application/json", + "content-type": "application/json", + "authorization": f"Bearer {self.runtime.credentials['siliconFlow_api_key']}", + } + + payload = { + "prompt": tool_parameters.get("prompt"), + "image_size": tool_parameters.get("image_size", "1024x1024"), + "seed": tool_parameters.get("seed"), + "num_inference_steps": tool_parameters.get("num_inference_steps", 20), + } + model = tool_parameters.get("model", "schnell") + url = FLUX_URL.get(model) + if model == "dev": + payload["model"] = "black-forest-labs/FLUX.1-dev" + + response = requests.post(url, json=payload, headers=headers) + if response.status_code != 200: + return self.create_text_message(f"Got Error Response:{response.text}") + + res = response.json() + result = [self.create_json_message(res)] + for image in res.get("images", []): + result.append(self.create_image_message(image=image.get("url"), save_as=self.VariableKey.IMAGE.value)) + return result diff --git a/api/core/tools/provider/builtin/siliconflow/tools/flux.yaml b/api/core/tools/provider/builtin/siliconflow/tools/flux.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d06b9bf3e1f489198a8e568ac4b247226619077c --- /dev/null +++ b/api/core/tools/provider/builtin/siliconflow/tools/flux.yaml @@ -0,0 +1,88 @@ +identity: + name: flux + author: hjlarry + label: + en_US: Flux + icon: icon.svg +description: + human: + en_US: Generate image via SiliconFlow's flux model. + llm: This tool is used to generate image from prompt via SiliconFlow's flux model. +parameters: + - name: prompt + type: string + required: true + label: + en_US: prompt + zh_Hans: 提示词 + human_description: + en_US: The text prompt used to generate the image. + zh_Hans: 建议用英文的生成图片提示词以获得更好的生成效果。 + llm_description: this prompt text will be used to generate image. + form: llm + - name: model + type: select + required: true + options: + - value: schnell + label: + en_US: Flux.1-schnell + - value: dev + label: + en_US: Flux.1-dev + default: schnell + label: + en_US: Choose Image Model + zh_Hans: 选择生成图片的模型 + form: form + - name: image_size + type: select + required: true + options: + - value: 1024x1024 + label: + en_US: 1024x1024 + - value: 768x1024 + label: + en_US: 768x1024 + - value: 576x1024 + label: + en_US: 576x1024 + - value: 512x1024 + label: + en_US: 512x1024 + - value: 1024x576 + label: + en_US: 1024x576 + - value: 768x512 + label: + en_US: 768x512 + default: 1024x1024 + label: + en_US: Choose Image Size + zh_Hans: 选择生成的图片大小 + form: form + - name: num_inference_steps + type: number + required: true + default: 20 + min: 1 + max: 100 + label: + en_US: Num Inference Steps + zh_Hans: 生成图片的步数 + form: form + human_description: + en_US: The number of inference steps to perform. More steps produce higher quality but take longer. + zh_Hans: 执行的推理步骤数量。更多的步骤可以产生更高质量的结果,但需要更长的时间。 + - name: seed + type: number + min: 0 + max: 9999999999 + label: + en_US: Seed + zh_Hans: 种子 + human_description: + en_US: The same seed and prompt can produce similar images. + zh_Hans: 相同的种子和提示可以产生相似的图像。 + form: form diff --git a/api/core/tools/provider/builtin/siliconflow/tools/stable_diffusion.py b/api/core/tools/provider/builtin/siliconflow/tools/stable_diffusion.py new file mode 100644 index 0000000000000000000000000000000000000000..db43790c06aaa6660da381b34f503f98ae86c03e --- /dev/null +++ b/api/core/tools/provider/builtin/siliconflow/tools/stable_diffusion.py @@ -0,0 +1,49 @@ +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +SILICONFLOW_API_URL = "https://api.siliconflow.cn/v1/image/generations" + +SD_MODELS = { + "sd_3": "stabilityai/stable-diffusion-3-medium", + "sd_xl": "stabilityai/stable-diffusion-xl-base-1.0", + "sd_3.5_large": "stabilityai/stable-diffusion-3-5-large", +} + + +class StableDiffusionTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + headers = { + "accept": "application/json", + "content-type": "application/json", + "authorization": f"Bearer {self.runtime.credentials['siliconFlow_api_key']}", + } + + model = tool_parameters.get("model", "sd_3") + sd_model = SD_MODELS.get(model) + + payload = { + "model": sd_model, + "prompt": tool_parameters.get("prompt"), + "negative_prompt": tool_parameters.get("negative_prompt", ""), + "image_size": tool_parameters.get("image_size", "1024x1024"), + "batch_size": tool_parameters.get("batch_size", 1), + "seed": tool_parameters.get("seed"), + "guidance_scale": tool_parameters.get("guidance_scale", 7.5), + "num_inference_steps": tool_parameters.get("num_inference_steps", 20), + } + + response = requests.post(SILICONFLOW_API_URL, json=payload, headers=headers) + if response.status_code != 200: + return self.create_text_message(f"Got Error Response:{response.text}") + + res = response.json() + result = [self.create_json_message(res)] + for image in res.get("images", []): + result.append(self.create_image_message(image=image.get("url"), save_as=self.VariableKey.IMAGE.value)) + return result diff --git a/api/core/tools/provider/builtin/siliconflow/tools/stable_diffusion.yaml b/api/core/tools/provider/builtin/siliconflow/tools/stable_diffusion.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b330c92e163a380cb9fe7ce37b263d9d2748632c --- /dev/null +++ b/api/core/tools/provider/builtin/siliconflow/tools/stable_diffusion.yaml @@ -0,0 +1,124 @@ +identity: + name: stable_diffusion + author: hjlarry + label: + en_US: Stable Diffusion + icon: icon.svg +description: + human: + en_US: Generate image via SiliconFlow's stable diffusion model. + llm: This tool is used to generate image from prompt via SiliconFlow's stable diffusion model. +parameters: + - name: prompt + type: string + required: true + label: + en_US: prompt + zh_Hans: 提示词 + human_description: + en_US: The text prompt used to generate the image. + zh_Hans: 用于生成图片的文字提示词 + llm_description: this prompt text will be used to generate image. + form: llm + - name: negative_prompt + type: string + label: + en_US: negative prompt + zh_Hans: 负面提示词 + human_description: + en_US: Describe what you don't want included in the image. + zh_Hans: 描述您不希望包含在图片中的内容。 + llm_description: Describe what you don't want included in the image. + form: llm + - name: model + type: select + required: true + options: + - value: sd_3 + label: + en_US: Stable Diffusion 3 + - value: sd_xl + label: + en_US: Stable Diffusion XL + - value: sd_3.5_large + label: + en_US: Stable Diffusion 3.5 Large + default: sd_3 + label: + en_US: Choose Image Model + zh_Hans: 选择生成图片的模型 + form: form + - name: image_size + type: select + required: true + options: + - value: 1024x1024 + label: + en_US: 1024x1024 + - value: 1024x2048 + label: + en_US: 1024x2048 + - value: 1152x2048 + label: + en_US: 1152x2048 + - value: 1536x1024 + label: + en_US: 1536x1024 + - value: 1536x2048 + label: + en_US: 1536x2048 + - value: 2048x1152 + label: + en_US: 2048x1152 + default: 1024x1024 + label: + en_US: Choose Image Size + zh_Hans: 选择生成图片的大小 + form: form + - name: batch_size + type: number + required: true + default: 1 + min: 1 + max: 4 + label: + en_US: Number Images + zh_Hans: 生成图片的数量 + form: form + - name: guidance_scale + type: number + required: true + default: 7.5 + min: 0 + max: 100 + label: + en_US: Guidance Scale + zh_Hans: 与提示词紧密性 + human_description: + en_US: Classifier Free Guidance. How close you want the model to stick to your prompt when looking for a related image to show you. + zh_Hans: 无分类器引导。您希望模型在寻找相关图片向您展示时,与您的提示保持多紧密的关联度。 + form: form + - name: num_inference_steps + type: number + required: true + default: 20 + min: 1 + max: 100 + label: + en_US: Num Inference Steps + zh_Hans: 生成图片的步数 + human_description: + en_US: The number of inference steps to perform. More steps produce higher quality but take longer. + zh_Hans: 执行的推理步骤数量。更多的步骤可以产生更高质量的结果,但需要更长的时间。 + form: form + - name: seed + type: number + min: 0 + max: 9999999999 + label: + en_US: Seed + zh_Hans: 种子 + human_description: + en_US: The same seed and prompt can produce similar images. + zh_Hans: 相同的种子和提示可以产生相似的图像。 + form: form diff --git a/api/core/tools/provider/builtin/slack/_assets/icon.svg b/api/core/tools/provider/builtin/slack/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..e43c2c47dc128ebd1d0048d37910573de5a9fcfd --- /dev/null +++ b/api/core/tools/provider/builtin/slack/_assets/icon.svg @@ -0,0 +1,22 @@ + + + Slack + + + + + + + diff --git a/api/core/tools/provider/builtin/slack/slack.py b/api/core/tools/provider/builtin/slack/slack.py new file mode 100644 index 0000000000000000000000000000000000000000..2de7911f63072aa724f7c6c9743468741f3909e4 --- /dev/null +++ b/api/core/tools/provider/builtin/slack/slack.py @@ -0,0 +1,8 @@ +from core.tools.provider.builtin.slack.tools.slack_webhook import SlackWebhookTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class SlackProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + SlackWebhookTool() + pass diff --git a/api/core/tools/provider/builtin/slack/slack.yaml b/api/core/tools/provider/builtin/slack/slack.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1070ffbf038a4081c44cc202138f9f2a6391d7ce --- /dev/null +++ b/api/core/tools/provider/builtin/slack/slack.yaml @@ -0,0 +1,16 @@ +identity: + author: Pan YANG + name: slack + label: + en_US: Slack + zh_Hans: Slack + pt_BR: Slack + description: + en_US: Slack Webhook + zh_Hans: Slack Webhook + pt_BR: Slack Webhook + icon: icon.svg + tags: + - social + - productivity +credentials_for_provider: diff --git a/api/core/tools/provider/builtin/slack/tools/slack_webhook.py b/api/core/tools/provider/builtin/slack/tools/slack_webhook.py new file mode 100644 index 0000000000000000000000000000000000000000..85e0de76755898aa18719b3145518345e0966e83 --- /dev/null +++ b/api/core/tools/provider/builtin/slack/tools/slack_webhook.py @@ -0,0 +1,46 @@ +from typing import Any, Union + +import httpx + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class SlackWebhookTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Incoming Webhooks + API Document: https://api.slack.com/messaging/webhooks + """ + + content = tool_parameters.get("content", "") + if not content: + return self.create_text_message("Invalid parameter content") + + webhook_url = tool_parameters.get("webhook_url", "") + + if not webhook_url.startswith("https://hooks.slack.com/"): + return self.create_text_message( + f"Invalid parameter webhook_url ${webhook_url}, not a valid Slack webhook URL" + ) + + headers = { + "Content-Type": "application/json", + } + params = {} + payload = { + "text": content, + } + + try: + res = httpx.post(webhook_url, headers=headers, params=params, json=payload) + if res.is_success: + return self.create_text_message("Text message was sent successfully") + else: + return self.create_text_message( + f"Failed to send the text message, status code: {res.status_code}, response: {res.text}" + ) + except Exception as e: + return self.create_text_message("Failed to send message through webhook. {}".format(e)) diff --git a/api/core/tools/provider/builtin/slack/tools/slack_webhook.yaml b/api/core/tools/provider/builtin/slack/tools/slack_webhook.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b838d743733ec990fcf40487cbfea54a19137d06 --- /dev/null +++ b/api/core/tools/provider/builtin/slack/tools/slack_webhook.yaml @@ -0,0 +1,40 @@ +identity: + name: slack_webhook + author: Pan YANG + label: + en_US: Incoming Webhook to send message + zh_Hans: 通过入站 Webhook 发送消息 + pt_BR: Incoming Webhook to send message + icon: icon.svg +description: + human: + en_US: Sending a message on Slack via the Incoming Webhook + zh_Hans: 通过入站 Webhook 在 Slack 上发送消息 + pt_BR: Sending a message on Slack via the Incoming Webhook + llm: A tool for sending messages to a chat on Slack. +parameters: + - name: webhook_url + type: string + required: true + label: + en_US: Slack Incoming Webhook url + zh_Hans: Slack 入站 Webhook 的 url + pt_BR: Slack Incoming Webhook url + human_description: + en_US: Slack Incoming Webhook url + zh_Hans: Slack 入站 Webhook 的 url + pt_BR: Slack Incoming Webhook url + form: form + - name: content + type: string + required: true + label: + en_US: content + zh_Hans: 消息内容 + pt_BR: content + human_description: + en_US: Content to sent to the channel or person. + zh_Hans: 消息内容文本 + pt_BR: Content to sent to the channel or person. + llm_description: Content of the message + form: llm diff --git a/api/core/tools/provider/builtin/slidespeak/_assets/icon.png b/api/core/tools/provider/builtin/slidespeak/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..4cac578330b15602fe79b7151956de358e706f3e Binary files /dev/null and b/api/core/tools/provider/builtin/slidespeak/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/slidespeak/slidespeak.py b/api/core/tools/provider/builtin/slidespeak/slidespeak.py new file mode 100644 index 0000000000000000000000000000000000000000..14c7c4880e892f6f92d6c867db60574d6bb9540c --- /dev/null +++ b/api/core/tools/provider/builtin/slidespeak/slidespeak.py @@ -0,0 +1,28 @@ +from typing import Any + +import requests +from yarl import URL + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class SlideSpeakProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + api_key = credentials.get("slidespeak_api_key") + base_url = credentials.get("base_url") + + if not api_key: + raise ToolProviderCredentialValidationError("API key is missing") + + if base_url: + base_url = str(URL(base_url) / "v1") + + headers = {"Content-Type": "application/json", "X-API-Key": api_key} + + test_task_id = "xxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx" + url = f"{base_url or 'https://api.slidespeak.co/api/v1'}/task_status/{test_task_id}" + + response = requests.get(url, headers=headers) + if response.status_code != 200: + raise ToolProviderCredentialValidationError("Invalid SlidePeak API key") diff --git a/api/core/tools/provider/builtin/slidespeak/slidespeak.yaml b/api/core/tools/provider/builtin/slidespeak/slidespeak.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9f6927f1bdcdf38674d61e6bc0e013a6da181b62 --- /dev/null +++ b/api/core/tools/provider/builtin/slidespeak/slidespeak.yaml @@ -0,0 +1,22 @@ +identity: + author: Kalo Chin + name: slidespeak + label: + en_US: SlideSpeak + zh_Hans: SlideSpeak + description: + en_US: Generate presentation slides using SlideSpeak API + zh_Hans: 使用 SlideSpeak API 生成演示幻灯片 + icon: icon.png + +credentials_for_provider: + slidespeak_api_key: + type: secret-input + required: true + label: + en_US: API Key + zh_Hans: API 密钥 + placeholder: + en_US: Enter your SlideSpeak API key + zh_Hans: 输入您的 SlideSpeak API 密钥 + url: https://app.slidespeak.co/settings/developer diff --git a/api/core/tools/provider/builtin/slidespeak/tools/slides_generator.py b/api/core/tools/provider/builtin/slidespeak/tools/slides_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..aa4ee63e9767c9191d1f526776b345bddc465357 --- /dev/null +++ b/api/core/tools/provider/builtin/slidespeak/tools/slides_generator.py @@ -0,0 +1,163 @@ +import asyncio +from dataclasses import asdict, dataclass +from enum import Enum +from typing import Any, Optional, Union + +import aiohttp +from pydantic import ConfigDict + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.tool.builtin_tool import BuiltinTool + + +class SlidesGeneratorTool(BuiltinTool): + """ + Tool for generating presentations using the SlideSpeak API. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + headers: Optional[dict[str, str]] = None + base_url: Optional[str] = None + timeout: Optional[aiohttp.ClientTimeout] = None + poll_interval: Optional[int] = None + + class TaskState(Enum): + FAILURE = "FAILURE" + REVOKED = "REVOKED" + SUCCESS = "SUCCESS" + PENDING = "PENDING" + RECEIVED = "RECEIVED" + STARTED = "STARTED" + + @dataclass + class PresentationRequest: + plain_text: str + length: Optional[int] = None + theme: Optional[str] = None + + async def _generate_presentation( + self, + session: aiohttp.ClientSession, + request: PresentationRequest, + ) -> dict[str, Any]: + """Generate a new presentation asynchronously""" + async with session.post( + f"{self.base_url}/presentation/generate", + headers=self.headers, + json=asdict(request), + timeout=self.timeout, + ) as response: + response.raise_for_status() + return await response.json() + + async def _get_task_status( + self, + session: aiohttp.ClientSession, + task_id: str, + ) -> dict[str, Any]: + """Get the status of a task asynchronously""" + async with session.get( + f"{self.base_url}/task_status/{task_id}", + headers=self.headers, + timeout=self.timeout, + ) as response: + response.raise_for_status() + return await response.json() + + async def _wait_for_completion( + self, + session: aiohttp.ClientSession, + task_id: str, + ) -> str: + """Wait for task completion and return download URL""" + while True: + status = await self._get_task_status(session, task_id) + task_status = self.TaskState(status["task_status"]) + if task_status == self.TaskState.SUCCESS: + return status["task_result"]["url"] + if task_status in [self.TaskState.FAILURE, self.TaskState.REVOKED]: + raise Exception(f"Task failed with status: {task_status.value}") + await asyncio.sleep(self.poll_interval) + + async def _generate_slides( + self, + plain_text: str, + length: Optional[int], + theme: Optional[str], + ) -> str: + """Generate slides and return the download URL""" + async with aiohttp.ClientSession() as session: + request = self.PresentationRequest( + plain_text=plain_text, + length=length, + theme=theme, + ) + result = await self._generate_presentation(session, request) + task_id = result["task_id"] + download_url = await self._wait_for_completion(session, task_id) + return download_url + + async def _fetch_presentation( + self, + session: aiohttp.ClientSession, + download_url: str, + ) -> bytes: + """Fetch the presentation file from the download URL""" + async with session.get(download_url, timeout=self.timeout) as response: + response.raise_for_status() + return await response.read() + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """Synchronous invoke method that runs asynchronous code""" + + async def async_invoke(): + # Extract parameters + plain_text = tool_parameters.get("plain_text", "") + length = tool_parameters.get("length") + theme = tool_parameters.get("theme") + + # Ensure runtime and credentials + if not self.runtime or not self.runtime.credentials: + raise ToolProviderCredentialValidationError("Tool runtime or credentials are missing") + + # Get API key from credentials + api_key = self.runtime.credentials.get("slidespeak_api_key") + if not api_key: + raise ToolProviderCredentialValidationError("SlideSpeak API key is missing") + + # Set configuration + self.headers = { + "Content-Type": "application/json", + "X-API-Key": api_key, + } + self.base_url = "https://api.slidespeak.co/api/v1" + self.timeout = aiohttp.ClientTimeout(total=30) + self.poll_interval = 2 + + # Run the asynchronous slide generation + try: + download_url = await self._generate_slides(plain_text, length, theme) + + # Fetch the presentation file + async with aiohttp.ClientSession() as session: + presentation_bytes = await self._fetch_presentation(session, download_url) + + return [ + self.create_text_message(download_url), + self.create_blob_message( + blob=presentation_bytes, + meta={"mime_type": "application/vnd.openxmlformats-officedocument.presentationml.presentation"}, + ), + ] + except Exception as e: + return [self.create_text_message(f"An error occurred: {str(e)}")] + + # Run the asynchronous code synchronously + result = asyncio.run(async_invoke()) + return result diff --git a/api/core/tools/provider/builtin/slidespeak/tools/slides_generator.yaml b/api/core/tools/provider/builtin/slidespeak/tools/slides_generator.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f881dadb20f82be4b862c3a09dac8e56997c3dec --- /dev/null +++ b/api/core/tools/provider/builtin/slidespeak/tools/slides_generator.yaml @@ -0,0 +1,102 @@ +identity: + name: slide_generator + author: Kalo Chin + label: + en_US: Slides Generator + zh_Hans: 幻灯片生成器 +description: + human: + en_US: Generate presentation slides from text using SlideSpeak API. + zh_Hans: 使用 SlideSpeak API 从文本生成演示幻灯片。 + llm: This tool converts text input into a presentation using the SlideSpeak API service, with options for slide length and theme. +parameters: + - name: plain_text + type: string + required: true + label: + en_US: Topic or Content + zh_Hans: 主题或内容 + human_description: + en_US: The topic or content to be converted into presentation slides. + zh_Hans: 需要转换为幻灯片的内容或主题。 + llm_description: A string containing the topic or content to be transformed into presentation slides. + form: llm + - name: length + type: number + required: false + label: + en_US: Number of Slides + zh_Hans: 幻灯片数量 + human_description: + en_US: The desired number of slides in the presentation (optional). + zh_Hans: 演示文稿中所需的幻灯片数量(可选)。 + llm_description: Optional parameter specifying the number of slides to generate. + form: form + - name: theme + type: select + required: false + label: + en_US: Presentation Theme + zh_Hans: 演示主题 + human_description: + en_US: The visual theme for the presentation (optional). + zh_Hans: 演示文稿的视觉主题(可选)。 + llm_description: Optional parameter specifying the presentation theme. + options: + - label: + en_US: Adam + zh_Hans: Adam + value: adam + - label: + en_US: Aurora + zh_Hans: Aurora + value: aurora + - label: + en_US: Bruno + zh_Hans: Bruno + value: bruno + - label: + en_US: Clyde + zh_Hans: Clyde + value: clyde + - label: + en_US: Daniel + zh_Hans: Daniel + value: daniel + - label: + en_US: Default + zh_Hans: Default + value: default + - label: + en_US: Eddy + zh_Hans: Eddy + value: eddy + - label: + en_US: Felix + zh_Hans: Felix + value: felix + - label: + en_US: Gradient + zh_Hans: Gradient + value: gradient + - label: + en_US: Iris + zh_Hans: Iris + value: iris + - label: + en_US: Lavender + zh_Hans: Lavender + value: lavender + - label: + en_US: Monolith + zh_Hans: Monolith + value: monolith + - label: + en_US: Nebula + zh_Hans: Nebula + value: nebula + - label: + en_US: Nexus + zh_Hans: Nexus + value: nexus + form: form diff --git a/api/core/tools/provider/builtin/spark/__init__.py b/api/core/tools/provider/builtin/spark/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/tools/provider/builtin/spark/_assets/icon.svg b/api/core/tools/provider/builtin/spark/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..ef0a9131a48e43a7968e53366da399b6dd931b8c --- /dev/null +++ b/api/core/tools/provider/builtin/spark/_assets/icon.svg @@ -0,0 +1,5 @@ + + + + + diff --git a/api/core/tools/provider/builtin/spark/spark.py b/api/core/tools/provider/builtin/spark/spark.py new file mode 100644 index 0000000000000000000000000000000000000000..e0b1a58a3f679adf8f433a251bbdda4fda45263b --- /dev/null +++ b/api/core/tools/provider/builtin/spark/spark.py @@ -0,0 +1,36 @@ +import json + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.spark.tools.spark_img_generation import spark_response +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class SparkProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + try: + if "APPID" not in credentials or not credentials.get("APPID"): + raise ToolProviderCredentialValidationError("APPID is required.") + if "APISecret" not in credentials or not credentials.get("APISecret"): + raise ToolProviderCredentialValidationError("APISecret is required.") + if "APIKey" not in credentials or not credentials.get("APIKey"): + raise ToolProviderCredentialValidationError("APIKey is required.") + + appid = credentials.get("APPID") + apisecret = credentials.get("APISecret") + apikey = credentials.get("APIKey") + prompt = "a cute black dog" + + try: + response = spark_response(prompt, appid, apikey, apisecret) + data = json.loads(response) + code = data["header"]["code"] + + if code == 0: + # 0 success, + pass + else: + raise ToolProviderCredentialValidationError("image generate error, code:{}".format(code)) + except Exception as e: + raise ToolProviderCredentialValidationError("APPID APISecret APIKey is invalid. {}".format(e)) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/spark/spark.yaml b/api/core/tools/provider/builtin/spark/spark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fa1543443a2af8908cee7d76ddad7851cf77a64d --- /dev/null +++ b/api/core/tools/provider/builtin/spark/spark.yaml @@ -0,0 +1,61 @@ +identity: + author: Onelevenvy + name: spark + label: + en_US: Spark + zh_Hans: 讯飞星火 + pt_BR: Spark + description: + en_US: Spark Platform Toolkit + zh_Hans: 讯飞星火平台工具 + pt_BR: Pacote de Ferramentas da Plataforma Spark + icon: icon.svg + tags: + - image +credentials_for_provider: + APPID: + type: secret-input + required: true + label: + en_US: Spark APPID + zh_Hans: APPID + pt_BR: Spark APPID + help: + en_US: Please input your APPID + zh_Hans: 请输入你的 APPID + pt_BR: Please input your APPID + placeholder: + en_US: Please input your APPID + zh_Hans: 请输入你的 APPID + pt_BR: Please input your APPID + APISecret: + type: secret-input + required: true + label: + en_US: Spark APISecret + zh_Hans: APISecret + pt_BR: Spark APISecret + help: + en_US: Please input your Spark APISecret + zh_Hans: 请输入你的 APISecret + pt_BR: Please input your Spark APISecret + placeholder: + en_US: Please input your Spark APISecret + zh_Hans: 请输入你的 APISecret + pt_BR: Please input your Spark APISecret + APIKey: + type: secret-input + required: true + label: + en_US: Spark APIKey + zh_Hans: APIKey + pt_BR: Spark APIKey + help: + en_US: Please input your Spark APIKey + zh_Hans: 请输入你的 APIKey + pt_BR: Please input your Spark APIKey + placeholder: + en_US: Please input your Spark APIKey + zh_Hans: 请输入你的 APIKey + pt_BR: Please input Spark APIKey + url: https://console.xfyun.cn/services diff --git a/api/core/tools/provider/builtin/spark/tools/spark_img_generation.py b/api/core/tools/provider/builtin/spark/tools/spark_img_generation.py new file mode 100644 index 0000000000000000000000000000000000000000..81d9e8d94185f745d4d698517ec4aee57582919c --- /dev/null +++ b/api/core/tools/provider/builtin/spark/tools/spark_img_generation.py @@ -0,0 +1,139 @@ +import base64 +import hashlib +import hmac +import json +from base64 import b64decode +from datetime import datetime +from time import mktime +from typing import Any, Union +from urllib.parse import urlencode +from wsgiref.handlers import format_date_time + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class AssembleHeaderError(Exception): + def __init__(self, msg): + self.message = msg + + +class Url: + def __init__(self, host, path, schema): + self.host = host + self.path = path + self.schema = schema + + +# calculate sha256 and encode to base64 +def sha256base64(data): + sha256 = hashlib.sha256() + sha256.update(data) + digest = base64.b64encode(sha256.digest()).decode(encoding="utf-8") + return digest + + +def parse_url(request_url): + stidx = request_url.index("://") + host = request_url[stidx + 3 :] + schema = request_url[: stidx + 3] + edidx = host.index("/") + if edidx <= 0: + raise AssembleHeaderError("invalid request url:" + request_url) + path = host[edidx:] + host = host[:edidx] + u = Url(host, path, schema) + return u + + +def assemble_ws_auth_url(request_url, method="GET", api_key="", api_secret=""): + u = parse_url(request_url) + host = u.host + path = u.path + now = datetime.now() + date = format_date_time(mktime(now.timetuple())) + signature_origin = "host: {}\ndate: {}\n{} {} HTTP/1.1".format(host, date, method, path) + signature_sha = hmac.new( + api_secret.encode("utf-8"), + signature_origin.encode("utf-8"), + digestmod=hashlib.sha256, + ).digest() + signature_sha = base64.b64encode(signature_sha).decode(encoding="utf-8") + authorization_origin = ( + f'api_key="{api_key}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature_sha}"' + ) + + authorization = base64.b64encode(authorization_origin.encode("utf-8")).decode(encoding="utf-8") + values = {"host": host, "date": date, "authorization": authorization} + + return request_url + "?" + urlencode(values) + + +def get_body(appid, text): + body = { + "header": {"app_id": appid, "uid": "123456789"}, + "parameter": {"chat": {"domain": "general", "temperature": 0.5, "max_tokens": 4096}}, + "payload": {"message": {"text": [{"role": "user", "content": text}]}}, + } + return body + + +def spark_response(text, appid, apikey, apisecret): + host = "http://spark-api.cn-huabei-1.xf-yun.com/v2.1/tti" + url = assemble_ws_auth_url(host, method="POST", api_key=apikey, api_secret=apisecret) + content = get_body(appid, text) + response = requests.post(url, json=content, headers={"content-type": "application/json"}).text + return response + + +class SparkImgGeneratorTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + + if "APPID" not in self.runtime.credentials or not self.runtime.credentials.get("APPID"): + return self.create_text_message("APPID is required.") + if "APISecret" not in self.runtime.credentials or not self.runtime.credentials.get("APISecret"): + return self.create_text_message("APISecret is required.") + if "APIKey" not in self.runtime.credentials or not self.runtime.credentials.get("APIKey"): + return self.create_text_message("APIKey is required.") + + prompt = tool_parameters.get("prompt", "") + if not prompt: + return self.create_text_message("Please input prompt") + res = self.img_generation(prompt) + result = [] + for image in res: + result.append( + self.create_blob_message( + blob=b64decode(image["base64_image"]), + meta={"mime_type": "image/png"}, + save_as=self.VariableKey.IMAGE.value, + ) + ) + return result + + def img_generation(self, prompt): + response = spark_response( + text=prompt, + appid=self.runtime.credentials.get("APPID"), + apikey=self.runtime.credentials.get("APIKey"), + apisecret=self.runtime.credentials.get("APISecret"), + ) + data = json.loads(response) + code = data["header"]["code"] + if code != 0: + return self.create_text_message(f"error: {code}, {data}") + else: + text = data["payload"]["choices"]["text"] + image_content = text[0] + image_base = image_content["content"] + json_data = {"base64_image": image_base} + return [json_data] diff --git a/api/core/tools/provider/builtin/spark/tools/spark_img_generation.yaml b/api/core/tools/provider/builtin/spark/tools/spark_img_generation.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d44bbc9564ef888d991a77a52f1e536033bf1d61 --- /dev/null +++ b/api/core/tools/provider/builtin/spark/tools/spark_img_generation.yaml @@ -0,0 +1,36 @@ +identity: + name: spark_img_generation + author: Onelevenvy + label: + en_US: Spark Image Generation + zh_Hans: 图片生成 + pt_BR: Geração de imagens Spark + icon: icon.svg + description: + en_US: Spark Image Generation + zh_Hans: 图片生成 + pt_BR: Geração de imagens Spark +description: + human: + en_US: Generate images based on user input, with image generation API + provided by Spark + zh_Hans: 根据用户的输入生成图片,由讯飞星火提供图片生成api + pt_BR: Gerar imagens com base na entrada do usuário, com API de geração + de imagem fornecida pela Spark + llm: spark_img_generation is a tool used to generate images from text +parameters: + - name: prompt + type: string + required: true + label: + en_US: Prompt + zh_Hans: 提示词 + pt_BR: Prompt + human_description: + en_US: Image prompt + zh_Hans: 图像提示词 + pt_BR: Image prompt + llm_description: Image prompt of spark_img_generation tooll, you should + describe the image you want to generate as a list of words as possible + as detailed + form: llm diff --git a/api/core/tools/provider/builtin/spider/_assets/icon.svg b/api/core/tools/provider/builtin/spider/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..604a09d01d744400bb218ec9fd9bab6728171206 --- /dev/null +++ b/api/core/tools/provider/builtin/spider/_assets/icon.svg @@ -0,0 +1 @@ +Spider v1 Logo diff --git a/api/core/tools/provider/builtin/spider/spider.py b/api/core/tools/provider/builtin/spider/spider.py new file mode 100644 index 0000000000000000000000000000000000000000..5959555318722ecc7dd703c6ab532baadde7baa8 --- /dev/null +++ b/api/core/tools/provider/builtin/spider/spider.py @@ -0,0 +1,20 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.spider.spiderApp import Spider +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class SpiderProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + app = Spider(api_key=credentials["spider_api_key"]) + app.scrape_url(url="https://spider.cloud") + except AttributeError as e: + # Handle cases where NoneType is not iterable, which might indicate API issues + if "NoneType" in str(e) and "not iterable" in str(e): + raise ToolProviderCredentialValidationError("API is currently down, try again in 15 minutes", str(e)) + else: + raise ToolProviderCredentialValidationError("An unexpected error occurred.", str(e)) + except Exception as e: + raise ToolProviderCredentialValidationError("An unexpected error occurred.", str(e)) diff --git a/api/core/tools/provider/builtin/spider/spider.yaml b/api/core/tools/provider/builtin/spider/spider.yaml new file mode 100644 index 0000000000000000000000000000000000000000..45702c85ddea24363cb7a3d65ed52ff760cc8f8a --- /dev/null +++ b/api/core/tools/provider/builtin/spider/spider.yaml @@ -0,0 +1,27 @@ +identity: + author: William Espegren + name: spider + label: + en_US: Spider + zh_CN: Spider + description: + en_US: Spider API integration, returning LLM-ready data by scraping & crawling websites. + zh_CN: Spider API 集成,通过爬取和抓取网站返回 LLM-ready 数据。 + icon: icon.svg + tags: + - search + - utilities +credentials_for_provider: + spider_api_key: + type: secret-input + required: true + label: + en_US: Spider API Key + zh_CN: Spider API 密钥 + placeholder: + en_US: Please input your Spider API key + zh_CN: 请输入您的 Spider API 密钥 + help: + en_US: Get your Spider API key from your Spider dashboard + zh_CN: 从您的 Spider 仪表板中获取 Spider API 密钥。 + url: https://spider.cloud/ diff --git a/api/core/tools/provider/builtin/spider/spiderApp.py b/api/core/tools/provider/builtin/spider/spiderApp.py new file mode 100644 index 0000000000000000000000000000000000000000..4bc446a1a092a3609a188cf9a20f292c63a3575d --- /dev/null +++ b/api/core/tools/provider/builtin/spider/spiderApp.py @@ -0,0 +1,221 @@ +import os +from typing import Literal, Optional, TypedDict + +import requests + + +class RequestParamsDict(TypedDict, total=False): + url: Optional[str] + request: Optional[Literal["http", "chrome", "smart"]] + limit: Optional[int] + return_format: Optional[Literal["raw", "markdown", "html2text", "text", "bytes"]] + tld: Optional[bool] + depth: Optional[int] + cache: Optional[bool] + budget: Optional[dict[str, int]] + locale: Optional[str] + cookies: Optional[str] + stealth: Optional[bool] + headers: Optional[dict[str, str]] + anti_bot: Optional[bool] + metadata: Optional[bool] + viewport: Optional[dict[str, int]] + encoding: Optional[str] + subdomains: Optional[bool] + user_agent: Optional[str] + store_data: Optional[bool] + gpt_config: Optional[list[str]] + fingerprint: Optional[bool] + storageless: Optional[bool] + readability: Optional[bool] + proxy_enabled: Optional[bool] + respect_robots: Optional[bool] + query_selector: Optional[str] + full_resources: Optional[bool] + request_timeout: Optional[int] + run_in_background: Optional[bool] + skip_config_checks: Optional[bool] + + +class Spider: + def __init__(self, api_key: Optional[str] = None): + """ + Initialize the Spider with an API key. + + :param api_key: A string of the API key for Spider. Defaults to the SPIDER_API_KEY environment variable. + :raises ValueError: If no API key is provided. + """ + self.api_key = api_key or os.getenv("SPIDER_API_KEY") + if self.api_key is None: + raise ValueError("No API key provided") + + def api_post( + self, + endpoint: str, + data: dict, + stream: bool, + content_type: str = "application/json", + ): + """ + Send a POST request to the specified API endpoint. + + :param endpoint: The API endpoint to which the POST request is sent. + :param data: The data (dictionary) to be sent in the POST request. + :param stream: Boolean indicating if the response should be streamed. + :return: The JSON response or the raw response stream if stream is True. + """ + headers = self._prepare_headers(content_type) + response = self._post_request(f"https://api.spider.cloud/v1/{endpoint}", data, headers, stream) + + if stream: + return response + elif response.status_code == 200: + return response.json() + else: + self._handle_error(response, f"post to {endpoint}") + + def api_get(self, endpoint: str, stream: bool, content_type: str = "application/json"): + """ + Send a GET request to the specified endpoint. + + :param endpoint: The API endpoint from which to retrieve data. + :return: The JSON decoded response. + """ + headers = self._prepare_headers(content_type) + response = self._get_request(f"https://api.spider.cloud/v1/{endpoint}", headers, stream) + if response.status_code == 200: + return response.json() + else: + self._handle_error(response, f"get from {endpoint}") + + def get_credits(self): + """ + Retrieve the account's remaining credits. + + :return: JSON response containing the number of credits left. + """ + return self.api_get("credits", stream=False) + + def scrape_url( + self, + url: str, + params: Optional[RequestParamsDict] = None, + stream: bool = False, + content_type: str = "application/json", + ): + """ + Scrape data from the specified URL. + + :param url: The URL from which to scrape data. + :param params: Optional dictionary of additional parameters for the scrape request. + :return: JSON response containing the scraping results. + """ + params = params or {} + + # Add { "return_format": "markdown" } to the params if not already present + if "return_format" not in params: + params["return_format"] = "markdown" + + # Set limit to 1 + params["limit"] = 1 + + return self.api_post("crawl", {"url": url, **(params or {})}, stream, content_type) + + def crawl_url( + self, + url: str, + params: Optional[RequestParamsDict] = None, + stream: bool = False, + content_type: str = "application/json", + ): + """ + Start crawling at the specified URL. + + :param url: The URL to begin crawling. + :param params: Optional dictionary with additional parameters to customize the crawl. + :param stream: Boolean indicating if the response should be streamed. Defaults to False. + :return: JSON response or the raw response stream if streaming enabled. + """ + params = params or {} + + # Add { "return_format": "markdown" } to the params if not already present + if "return_format" not in params: + params["return_format"] = "markdown" + + return self.api_post("crawl", {"url": url, **(params or {})}, stream, content_type) + + def links( + self, + url: str, + params: Optional[RequestParamsDict] = None, + stream: bool = False, + content_type: str = "application/json", + ): + """ + Retrieve links from the specified URL. + + :param url: The URL from which to extract links. + :param params: Optional parameters for the link retrieval request. + :return: JSON response containing the links. + """ + return self.api_post("links", {"url": url, **(params or {})}, stream, content_type) + + def extract_contacts( + self, + url: str, + params: Optional[RequestParamsDict] = None, + stream: bool = False, + content_type: str = "application/json", + ): + """ + Extract contact information from the specified URL. + + :param url: The URL from which to extract contact information. + :param params: Optional parameters for the contact extraction. + :return: JSON response containing extracted contact details. + """ + return self.api_post( + "pipeline/extract-contacts", + {"url": url, **(params or {})}, + stream, + content_type, + ) + + def label( + self, + url: str, + params: Optional[RequestParamsDict] = None, + stream: bool = False, + content_type: str = "application/json", + ): + """ + Apply labeling to data extracted from the specified URL. + + :param url: The URL to label data from. + :param params: Optional parameters to guide the labeling process. + :return: JSON response with labeled data. + """ + return self.api_post("pipeline/label", {"url": url, **(params or {})}, stream, content_type) + + def _prepare_headers(self, content_type: str = "application/json"): + return { + "Content-Type": content_type, + "Authorization": f"Bearer {self.api_key}", + "User-Agent": "Spider-Client/0.0.27", + } + + def _post_request(self, url: str, data, headers, stream=False): + return requests.post(url, headers=headers, json=data, stream=stream) + + def _get_request(self, url: str, headers, stream=False): + return requests.get(url, headers=headers, stream=stream) + + def _delete_request(self, url: str, headers, stream=False): + return requests.delete(url, headers=headers, stream=stream) + + def _handle_error(self, response, action): + if response.status_code in {402, 409, 500}: + error_message = response.json().get("error", "Unknown error occurred") + raise Exception(f"Failed to {action}. Status code: {response.status_code}. Error: {error_message}") + else: + raise Exception(f"Unexpected error occurred while trying to {action}. Status code: {response.status_code}") diff --git a/api/core/tools/provider/builtin/spider/tools/scraper_crawler.py b/api/core/tools/provider/builtin/spider/tools/scraper_crawler.py new file mode 100644 index 0000000000000000000000000000000000000000..20d2daef550de1a6f3a2d7ce0d401d6b8dbc7bab --- /dev/null +++ b/api/core/tools/provider/builtin/spider/tools/scraper_crawler.py @@ -0,0 +1,49 @@ +from typing import Any, Union + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.spider.spiderApp import Spider +from core.tools.tool.builtin_tool import BuiltinTool + + +class ScrapeTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + # initialize the app object with the api key + app = Spider(api_key=self.runtime.credentials["spider_api_key"]) + + url = tool_parameters["url"] + mode = tool_parameters["mode"] + + options = { + "limit": tool_parameters.get("limit", 0), + "depth": tool_parameters.get("depth", 0), + "blacklist": tool_parameters.get("blacklist", "").split(",") if tool_parameters.get("blacklist") else [], + "whitelist": tool_parameters.get("whitelist", "").split(",") if tool_parameters.get("whitelist") else [], + "readability": tool_parameters.get("readability", False), + } + + result = "" + + try: + if mode == "scrape": + scrape_result = app.scrape_url( + url=url, + params=options, + ) + + for i in scrape_result: + result += "URL: " + i.get("url", "") + "\n" + result += "CONTENT: " + i.get("content", "") + "\n\n" + elif mode == "crawl": + crawl_result = app.crawl_url( + url=tool_parameters["url"], + params=options, + ) + for i in crawl_result: + result += "URL: " + i.get("url", "") + "\n" + result += "CONTENT: " + i.get("content", "") + "\n\n" + except Exception as e: + return self.create_text_message("An error occurred", str(e)) + + return self.create_text_message(result) diff --git a/api/core/tools/provider/builtin/spider/tools/scraper_crawler.yaml b/api/core/tools/provider/builtin/spider/tools/scraper_crawler.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5b20c2fc2f70ad5991286d6b17fa384be901971d --- /dev/null +++ b/api/core/tools/provider/builtin/spider/tools/scraper_crawler.yaml @@ -0,0 +1,102 @@ +identity: + name: scraper_crawler + author: William Espegren + label: + en_US: Web Scraper & Crawler + zh_Hans: 网页抓取与爬虫 +description: + human: + en_US: A tool for scraping & crawling webpages. Input should be a url. + zh_Hans: 用于抓取和爬取网页的工具。输入应该是一个网址。 + llm: A tool for scraping & crawling webpages. Input should be a url. +parameters: + - name: url + type: string + required: true + label: + en_US: URL + zh_Hans: 网址 + human_description: + en_US: url to be scraped or crawled + zh_Hans: 要抓取或爬取的网址 + llm_description: url to either be scraped or crawled + form: llm + - name: mode + type: select + required: true + options: + - value: scrape + label: + en_US: scrape + zh_Hans: 抓取 + - value: crawl + label: + en_US: crawl + zh_Hans: 爬取 + default: crawl + label: + en_US: Mode + zh_Hans: 模式 + human_description: + en_US: used for selecting to either scrape the website or crawl the entire website following subpages + zh_Hans: 用于选择抓取网站或爬取整个网站及其子页面 + form: form + - name: limit + type: number + required: false + label: + en_US: maximum number of pages to crawl + zh_Hans: 最大爬取页面数 + human_description: + en_US: specify the maximum number of pages to crawl per website. the crawler will stop after reaching this limit. + zh_Hans: 指定每个网站要爬取的最大页面数。爬虫将在达到此限制后停止。 + form: form + min: 0 + default: 0 + - name: depth + type: number + required: false + label: + en_US: maximum depth of pages to crawl + zh_Hans: 最大爬取深度 + human_description: + en_US: the crawl limit for maximum depth. + zh_Hans: 最大爬取深度的限制。 + form: form + min: 0 + default: 0 + - name: blacklist + type: string + required: false + label: + en_US: url patterns to exclude + zh_Hans: 要排除的URL模式 + human_description: + en_US: blacklist a set of paths that you do not want to crawl. you can use regex patterns to help with the list. + zh_Hans: 指定一组不想爬取的路径。您可以使用正则表达式模式来帮助定义列表。 + placeholder: + en_US: /blog/*, /about + form: form + - name: whitelist + type: string + required: false + label: + en_US: URL patterns to include + zh_Hans: 要包含的URL模式 + human_description: + en_US: Whitelist a set of paths that you want to crawl, ignoring all other routes that do not match the patterns. You can use regex patterns to help with the list. + zh_Hans: 指定一组要爬取的路径,忽略所有不匹配模式的其他路由。您可以使用正则表达式模式来帮助定义列表。 + placeholder: + en_US: /blog/*, /about + form: form + - name: readability + type: boolean + required: false + label: + en_US: Pre-process the content for LLM usage + zh_Hans: 仅返回页面的主要内容 + human_description: + en_US: Use Mozilla's readability to pre-process the content for reading. This may drastically improve the content for LLM usage. + zh_Hans: 如果启用,爬虫将仅返回页面的主要内容,不包括标题、导航、页脚等。 + form: form + default: false diff --git a/api/core/tools/provider/builtin/stability/_assets/icon.svg b/api/core/tools/provider/builtin/stability/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..56357a35557ac313dc252e99538ac600ef2b49c0 --- /dev/null +++ b/api/core/tools/provider/builtin/stability/_assets/icon.svg @@ -0,0 +1,10 @@ + + + + + + + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/stability/stability.py b/api/core/tools/provider/builtin/stability/stability.py new file mode 100644 index 0000000000000000000000000000000000000000..f09d81ac270288ba0d42d983aa02267dc2bf907d --- /dev/null +++ b/api/core/tools/provider/builtin/stability/stability.py @@ -0,0 +1,16 @@ +from typing import Any + +from core.tools.provider.builtin.stability.tools.base import BaseStabilityAuthorization +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class StabilityToolProvider(BuiltinToolProviderController, BaseStabilityAuthorization): + """ + This class is responsible for providing the stability tool. + """ + + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + """ + This method is responsible for validating the credentials. + """ + self.sd_validate_credentials(credentials) diff --git a/api/core/tools/provider/builtin/stability/stability.yaml b/api/core/tools/provider/builtin/stability/stability.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c3e01c1e314d51806609a457eef33287f0a93abc --- /dev/null +++ b/api/core/tools/provider/builtin/stability/stability.yaml @@ -0,0 +1,31 @@ +identity: + author: Dify + name: stability + label: + en_US: Stability + zh_Hans: Stability + pt_BR: Stability + description: + en_US: Activating humanity's potential through generative AI + zh_Hans: 通过生成式 AI 激活人类的潜力 + pt_BR: Activating humanity's potential through generative AI + icon: icon.svg + tags: + - image +credentials_for_provider: + api_key: + type: secret-input + required: true + label: + en_US: API key + zh_Hans: API key + pt_BR: API key + placeholder: + en_US: Please input your API key + zh_Hans: 请输入你的 API key + pt_BR: Please input your API key + help: + en_US: Get your API key from Stability + zh_Hans: 从 Stability 获取你的 API key + pt_BR: Get your API key from Stability + url: https://platform.stability.ai/account/keys diff --git a/api/core/tools/provider/builtin/stability/tools/base.py b/api/core/tools/provider/builtin/stability/tools/base.py new file mode 100644 index 0000000000000000000000000000000000000000..2d1cd928703151ec799cc2fafd0a114e60da4336 --- /dev/null +++ b/api/core/tools/provider/builtin/stability/tools/base.py @@ -0,0 +1,31 @@ +import requests +from yarl import URL + +from core.tools.errors import ToolProviderCredentialValidationError + + +class BaseStabilityAuthorization: + def sd_validate_credentials(self, credentials: dict): + """ + This method is responsible for validating the credentials. + """ + api_key = credentials.get("api_key", "") + if not api_key: + raise ToolProviderCredentialValidationError("API key is required.") + + response = requests.get( + URL("https://api.stability.ai") / "v1" / "user" / "account", + headers=self.generate_authorization_headers(credentials), + timeout=(5, 30), + ) + + if not response.ok: + raise ToolProviderCredentialValidationError("Invalid API key.") + + return True + + def generate_authorization_headers(self, credentials: dict) -> dict[str, str]: + """ + This method is responsible for generating the authorization headers. + """ + return {"Authorization": f"Bearer {credentials.get('api_key', '')}"} diff --git a/api/core/tools/provider/builtin/stability/tools/text2image.py b/api/core/tools/provider/builtin/stability/tools/text2image.py new file mode 100644 index 0000000000000000000000000000000000000000..6bcf315484ad509d365974dd486c91aa96ca5126 --- /dev/null +++ b/api/core/tools/provider/builtin/stability/tools/text2image.py @@ -0,0 +1,56 @@ +from typing import Any + +from httpx import post + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.provider.builtin.stability.tools.base import BaseStabilityAuthorization +from core.tools.tool.builtin_tool import BuiltinTool + + +class StableDiffusionTool(BuiltinTool, BaseStabilityAuthorization): + """ + This class is responsible for providing the stable diffusion tool. + """ + + model_endpoint_map: dict[str, str] = { + "sd3": "https://api.stability.ai/v2beta/stable-image/generate/sd3", + "sd3-turbo": "https://api.stability.ai/v2beta/stable-image/generate/sd3", + "core": "https://api.stability.ai/v2beta/stable-image/generate/core", + } + + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage | list[ToolInvokeMessage]: + """ + Invoke the tool. + """ + payload = { + "prompt": tool_parameters.get("prompt", ""), + "aspect_ratio": tool_parameters.get("aspect_ratio", "16:9") or tool_parameters.get("aspect_radio", "16:9"), + "mode": "text-to-image", + "seed": tool_parameters.get("seed", 0), + "output_format": "png", + } + + model = tool_parameters.get("model", "core") + + if model in {"sd3", "sd3-turbo"}: + payload["model"] = tool_parameters.get("model") + + if model != "sd3-turbo": + payload["negative_prompt"] = tool_parameters.get("negative_prompt", "") + + response = post( + self.model_endpoint_map[tool_parameters.get("model", "core")], + headers={ + "accept": "image/*", + **self.generate_authorization_headers(self.runtime.credentials), + }, + files={key: (None, str(value)) for key, value in payload.items()}, + timeout=(5, 30), + ) + + if not response.status_code == 200: + raise Exception(response.text) + + return self.create_blob_message( + blob=response.content, meta={"mime_type": "image/png"}, save_as=self.VariableKey.IMAGE.value + ) diff --git a/api/core/tools/provider/builtin/stability/tools/text2image.yaml b/api/core/tools/provider/builtin/stability/tools/text2image.yaml new file mode 100644 index 0000000000000000000000000000000000000000..21345f9f187f07fd5eb9aa0f24bd3f20c3a3ee8b --- /dev/null +++ b/api/core/tools/provider/builtin/stability/tools/text2image.yaml @@ -0,0 +1,142 @@ +identity: + name: stability_text2image + author: Dify + label: + en_US: StableDiffusion + zh_Hans: 稳定扩散 + pt_BR: StableDiffusion +description: + human: + en_US: A tool for generate images based on the text input + zh_Hans: 一个基于文本输入生成图像的工具 + pt_BR: A tool for generate images based on the text input + llm: A tool for generate images based on the text input +parameters: + - name: prompt + type: string + required: true + label: + en_US: Prompt + zh_Hans: 提示词 + pt_BR: Prompt + human_description: + en_US: used for generating images + zh_Hans: 用于生成图像 + pt_BR: used for generating images + llm_description: key words for generating images + form: llm + - name: model + type: select + default: sd3-turbo + required: true + label: + en_US: Model + zh_Hans: 模型 + pt_BR: Model + options: + - value: core + label: + en_US: Core + zh_Hans: Core + pt_BR: Core + - value: sd3 + label: + en_US: Stable Diffusion 3 + zh_Hans: Stable Diffusion 3 + pt_BR: Stable Diffusion 3 + - value: sd3-turbo + label: + en_US: Stable Diffusion 3 Turbo + zh_Hans: Stable Diffusion 3 Turbo + pt_BR: Stable Diffusion 3 Turbo + human_description: + en_US: Model for generating images + zh_Hans: 用于生成图像的模型 + pt_BR: Model for generating images + llm_description: Model for generating images + form: form + - name: negative_prompt + type: string + default: bad art, ugly, deformed, watermark, duplicated, discontinuous lines + required: false + label: + en_US: Negative Prompt + zh_Hans: 负面提示 + pt_BR: Negative Prompt + human_description: + en_US: Negative Prompt + zh_Hans: 负面提示 + pt_BR: Negative Prompt + llm_description: Negative Prompt + form: form + - name: seeds + type: number + default: 0 + required: false + label: + en_US: Seeds + zh_Hans: 种子 + pt_BR: Seeds + human_description: + en_US: Seeds + zh_Hans: 种子 + pt_BR: Seeds + llm_description: Seeds + min: 0 + max: 4294967294 + form: form + - name: aspect_ratio + type: select + default: '16:9' + options: + - value: '16:9' + label: + en_US: '16:9' + zh_Hans: '16:9' + pt_BR: '16:9' + - value: '1:1' + label: + en_US: '1:1' + zh_Hans: '1:1' + pt_BR: '1:1' + - value: '21:9' + label: + en_US: '21:9' + zh_Hans: '21:9' + pt_BR: '21:9' + - value: '2:3' + label: + en_US: '2:3' + zh_Hans: '2:3' + pt_BR: '2:3' + - value: '4:5' + label: + en_US: '4:5' + zh_Hans: '4:5' + pt_BR: '4:5' + - value: '5:4' + label: + en_US: '5:4' + zh_Hans: '5:4' + pt_BR: '5:4' + - value: '9:16' + label: + en_US: '9:16' + zh_Hans: '9:16' + pt_BR: '9:16' + - value: '9:21' + label: + en_US: '9:21' + zh_Hans: '9:21' + pt_BR: '9:21' + required: false + label: + en_US: Aspect Ratio + zh_Hans: 长宽比 + pt_BR: Aspect Ratio + human_description: + en_US: Aspect Ratio + zh_Hans: 长宽比 + pt_BR: Aspect Ratio + llm_description: Aspect Ratio + form: form diff --git a/api/core/tools/provider/builtin/stablediffusion/_assets/icon.png b/api/core/tools/provider/builtin/stablediffusion/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..fc372b28f1ccfd7bea27dfe7ef0450e98a0be7e1 Binary files /dev/null and b/api/core/tools/provider/builtin/stablediffusion/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/stablediffusion/stablediffusion.py b/api/core/tools/provider/builtin/stablediffusion/stablediffusion.py new file mode 100644 index 0000000000000000000000000000000000000000..abaa297cf36eb121b89a2b05d31807f30e07187b --- /dev/null +++ b/api/core/tools/provider/builtin/stablediffusion/stablediffusion.py @@ -0,0 +1,17 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.stablediffusion.tools.stable_diffusion import StableDiffusionTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class StableDiffusionProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + StableDiffusionTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).validate_models() + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/stablediffusion/stablediffusion.yaml b/api/core/tools/provider/builtin/stablediffusion/stablediffusion.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9b3c804f722dfc586fe9829cfb98d17a729d8ab0 --- /dev/null +++ b/api/core/tools/provider/builtin/stablediffusion/stablediffusion.yaml @@ -0,0 +1,42 @@ +identity: + author: Dify + name: stablediffusion + label: + en_US: Stable Diffusion + zh_Hans: Stable Diffusion + pt_BR: Stable Diffusion + description: + en_US: Stable Diffusion is a tool for generating images which can be deployed locally. + zh_Hans: Stable Diffusion 是一个可以在本地部署的图片生成的工具。 + pt_BR: Stable Diffusion is a tool for generating images which can be deployed locally. + icon: icon.png + tags: + - image +credentials_for_provider: + base_url: + type: secret-input + required: true + label: + en_US: Base URL + zh_Hans: StableDiffusion服务器的Base URL + pt_BR: Base URL + placeholder: + en_US: Please input your StableDiffusion server's Base URL + zh_Hans: 请输入你的 StableDiffusion 服务器的 Base URL + pt_BR: Please input your StableDiffusion server's Base URL + model: + type: text-input + required: true + label: + en_US: Model + zh_Hans: 模型 + pt_BR: Model + placeholder: + en_US: Please input your model + zh_Hans: 请输入你的模型名称 + pt_BR: Please input your model + help: + en_US: The model name of the StableDiffusion server + zh_Hans: StableDiffusion服务器的模型名称 + pt_BR: The model name of the StableDiffusion server + url: https://docs.dify.ai/tutorials/tool-configuration/stable-diffusion diff --git a/api/core/tools/provider/builtin/stablediffusion/tools/stable_diffusion.py b/api/core/tools/provider/builtin/stablediffusion/tools/stable_diffusion.py new file mode 100644 index 0000000000000000000000000000000000000000..64fdc961b4c5dbf2671415b46182e56a0823791a --- /dev/null +++ b/api/core/tools/provider/builtin/stablediffusion/tools/stable_diffusion.py @@ -0,0 +1,390 @@ +import io +import json +from base64 import b64decode, b64encode +from copy import deepcopy +from typing import Any, Union + +from httpx import get, post +from PIL import Image +from yarl import URL + +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter, ToolParameterOption +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.tool.builtin_tool import BuiltinTool + +# All commented out parameters default to null +DRAW_TEXT_OPTIONS = { + # Prompts + "prompt": "", + "negative_prompt": "", + # "styles": [], + # Seeds + "seed": -1, + "subseed": -1, + "subseed_strength": 0, + "seed_resize_from_h": -1, + "seed_resize_from_w": -1, + # Samplers + "sampler_name": "DPM++ 2M", + # "scheduler": "", + # "sampler_index": "Automatic", + # Latent Space Options + "batch_size": 1, + "n_iter": 1, + "steps": 10, + "cfg_scale": 7, + "width": 512, + "height": 512, + # "restore_faces": True, + # "tiling": True, + "do_not_save_samples": False, + "do_not_save_grid": False, + # "eta": 0, + # "denoising_strength": 0.75, + # "s_min_uncond": 0, + # "s_churn": 0, + # "s_tmax": 0, + # "s_tmin": 0, + # "s_noise": 0, + "override_settings": {}, + "override_settings_restore_afterwards": True, + # Refinement Options + "refiner_checkpoint": "", + "refiner_switch_at": 0, + "disable_extra_networks": False, + # "firstpass_image": "", + # "comments": "", + # High-Resolution Options + "enable_hr": False, + "firstphase_width": 0, + "firstphase_height": 0, + "hr_scale": 2, + # "hr_upscaler": "", + "hr_second_pass_steps": 0, + "hr_resize_x": 0, + "hr_resize_y": 0, + # "hr_checkpoint_name": "", + # "hr_sampler_name": "", + # "hr_scheduler": "", + "hr_prompt": "", + "hr_negative_prompt": "", + # Task Options + # "force_task_id": "", + # Script Options + # "script_name": "", + "script_args": [], + # Output Options + "send_images": True, + "save_images": False, + "alwayson_scripts": {}, + # "infotext": "", +} + + +class StableDiffusionTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + # base url + base_url = self.runtime.credentials.get("base_url", None) + if not base_url: + return self.create_text_message("Please input base_url") + + if tool_parameters.get("model"): + self.runtime.credentials["model"] = tool_parameters["model"] + + model = self.runtime.credentials.get("model", None) + if not model: + return self.create_text_message("Please input model") + + # set model + try: + url = str(URL(base_url) / "sdapi" / "v1" / "options") + response = post(url, data=json.dumps({"sd_model_checkpoint": model})) + if response.status_code != 200: + raise ToolProviderCredentialValidationError("Failed to set model, please tell user to set model") + except Exception as e: + raise ToolProviderCredentialValidationError("Failed to set model, please tell user to set model") + + # get image id and image variable + image_id = tool_parameters.get("image_id", "") + image_variable = self.get_default_image_variable() + # Return text2img if there's no image ID or no image variable + if not image_id or not image_variable: + return self.text2img(base_url=base_url, tool_parameters=tool_parameters) + + # Proceed with image-to-image generation + return self.img2img(base_url=base_url, tool_parameters=tool_parameters) + + def validate_models(self) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + validate models + """ + try: + base_url = self.runtime.credentials.get("base_url", None) + if not base_url: + raise ToolProviderCredentialValidationError("Please input base_url") + model = self.runtime.credentials.get("model", None) + if not model: + raise ToolProviderCredentialValidationError("Please input model") + + api_url = str(URL(base_url) / "sdapi" / "v1" / "sd-models") + response = get(url=api_url, timeout=10) + if response.status_code == 404: + # try draw a picture + self._invoke( + user_id="test", + tool_parameters={ + "prompt": "a cat", + "width": 1024, + "height": 1024, + "steps": 1, + "lora": "", + }, + ) + elif response.status_code != 200: + raise ToolProviderCredentialValidationError("Failed to get models") + else: + models = [d["model_name"] for d in response.json()] + if len([d for d in models if d == model]) > 0: + return self.create_text_message(json.dumps(models)) + else: + raise ToolProviderCredentialValidationError(f"model {model} does not exist") + except Exception as e: + raise ToolProviderCredentialValidationError(f"Failed to get models, {e}") + + def get_sd_models(self) -> list[str]: + """ + get sd models + """ + try: + base_url = self.runtime.credentials.get("base_url", None) + if not base_url: + return [] + api_url = str(URL(base_url) / "sdapi" / "v1" / "sd-models") + response = get(url=api_url, timeout=(2, 10)) + if response.status_code != 200: + return [] + else: + return [d["model_name"] for d in response.json()] + except Exception as e: + return [] + + def get_sample_methods(self) -> list[str]: + """ + get sample method + """ + try: + base_url = self.runtime.credentials.get("base_url", None) + if not base_url: + return [] + api_url = str(URL(base_url) / "sdapi" / "v1" / "samplers") + response = get(url=api_url, timeout=(2, 10)) + if response.status_code != 200: + return [] + else: + return [d["name"] for d in response.json()] + except Exception as e: + return [] + + def img2img( + self, base_url: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + generate image + """ + + # Fetch the binary data of the image + image_variable = self.get_default_image_variable() + image_binary = self.get_variable_file(image_variable.name) + if not image_binary: + return self.create_text_message("Image not found, please request user to generate image firstly.") + + # Convert image to RGB and save as PNG + try: + with Image.open(io.BytesIO(image_binary)) as image, io.BytesIO() as buffer: + image.convert("RGB").save(buffer, format="PNG") + image_binary = buffer.getvalue() + except Exception as e: + return self.create_text_message(f"Failed to process the image: {str(e)}") + + # copy draw options + draw_options = deepcopy(DRAW_TEXT_OPTIONS) + # set image options + model = tool_parameters.get("model", "") + draw_options_image = { + "init_images": [b64encode(image_binary).decode("utf-8")], + "denoising_strength": 0.9, + "restore_faces": False, + "script_args": [], + "override_settings": {"sd_model_checkpoint": model}, + "resize_mode": 0, + "image_cfg_scale": 0, + # "mask": None, + "mask_blur_x": 4, + "mask_blur_y": 4, + "mask_blur": 0, + "mask_round": True, + "inpainting_fill": 0, + "inpaint_full_res": True, + "inpaint_full_res_padding": 0, + "inpainting_mask_invert": 0, + "initial_noise_multiplier": 0, + # "latent_mask": None, + "include_init_images": True, + } + # update key and values + draw_options.update(draw_options_image) + draw_options.update(tool_parameters) + + # get prompt lora model + prompt = tool_parameters.get("prompt", "") + lora = tool_parameters.get("lora", "") + model = tool_parameters.get("model", "") + if lora: + draw_options["prompt"] = f"{lora},{prompt}" + else: + draw_options["prompt"] = prompt + + try: + url = str(URL(base_url) / "sdapi" / "v1" / "img2img") + response = post(url, data=json.dumps(draw_options), timeout=120) + if response.status_code != 200: + return self.create_text_message("Failed to generate image") + + image = response.json()["images"][0] + + return self.create_blob_message( + blob=b64decode(image), meta={"mime_type": "image/png"}, save_as=self.VariableKey.IMAGE.value + ) + + except Exception as e: + return self.create_text_message("Failed to generate image") + + def text2img( + self, base_url: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + generate image + """ + # copy draw options + draw_options = deepcopy(DRAW_TEXT_OPTIONS) + draw_options.update(tool_parameters) + # get prompt lora model + prompt = tool_parameters.get("prompt", "") + lora = tool_parameters.get("lora", "") + model = tool_parameters.get("model", "") + if lora: + draw_options["prompt"] = f"{lora},{prompt}" + else: + draw_options["prompt"] = prompt + draw_options["override_settings"]["sd_model_checkpoint"] = model + + try: + url = str(URL(base_url) / "sdapi" / "v1" / "txt2img") + response = post(url, data=json.dumps(draw_options), timeout=120) + if response.status_code != 200: + return self.create_text_message("Failed to generate image") + + image = response.json()["images"][0] + + return self.create_blob_message( + blob=b64decode(image), meta={"mime_type": "image/png"}, save_as=self.VariableKey.IMAGE.value + ) + + except Exception as e: + return self.create_text_message("Failed to generate image") + + def get_runtime_parameters(self) -> list[ToolParameter]: + parameters = [ + ToolParameter( + name="prompt", + label=I18nObject(en_US="Prompt", zh_Hans="Prompt"), + human_description=I18nObject( + en_US="Image prompt, you can check the official documentation of Stable Diffusion", + zh_Hans="图像提示词,您可以查看 Stable Diffusion 的官方文档", + ), + type=ToolParameter.ToolParameterType.STRING, + form=ToolParameter.ToolParameterForm.LLM, + llm_description="Image prompt of Stable Diffusion, you should describe the image you want to generate" + " as a list of words as possible as detailed, the prompt must be written in English.", + required=True, + ), + ] + if len(self.list_default_image_variables()) != 0: + parameters.append( + ToolParameter( + name="image_id", + label=I18nObject(en_US="image_id", zh_Hans="image_id"), + human_description=I18nObject( + en_US="Image id of the image you want to generate based on, if you want to generate image based" + " on the default image, you can leave this field empty.", + zh_Hans="您想要生成的图像的图像 ID,如果您想要基于默认图像生成图像,则可以将此字段留空。", + ), + type=ToolParameter.ToolParameterType.STRING, + form=ToolParameter.ToolParameterForm.LLM, + llm_description="Image id of the original image, you can leave this field empty if you want to" + " generate a new image.", + required=True, + options=[ + ToolParameterOption(value=i.name, label=I18nObject(en_US=i.name, zh_Hans=i.name)) + for i in self.list_default_image_variables() + ], + ) + ) + + if self.runtime.credentials: + try: + models = self.get_sd_models() + if len(models) != 0: + parameters.append( + ToolParameter( + name="model", + label=I18nObject(en_US="Model", zh_Hans="Model"), + human_description=I18nObject( + en_US="Model of Stable Diffusion, you can check the official documentation" + " of Stable Diffusion", + zh_Hans="Stable Diffusion 的模型,您可以查看 Stable Diffusion 的官方文档", + ), + type=ToolParameter.ToolParameterType.SELECT, + form=ToolParameter.ToolParameterForm.FORM, + llm_description="Model of Stable Diffusion, you can check the official documentation" + " of Stable Diffusion", + required=True, + default=models[0], + options=[ + ToolParameterOption(value=i, label=I18nObject(en_US=i, zh_Hans=i)) for i in models + ], + ) + ) + + except: + pass + + sample_methods = self.get_sample_methods() + if len(sample_methods) != 0: + parameters.append( + ToolParameter( + name="sampler_name", + label=I18nObject(en_US="Sampling method", zh_Hans="Sampling method"), + human_description=I18nObject( + en_US="Sampling method of Stable Diffusion, you can check the official documentation" + " of Stable Diffusion", + zh_Hans="Stable Diffusion 的Sampling method,您可以查看 Stable Diffusion 的官方文档", + ), + type=ToolParameter.ToolParameterType.SELECT, + form=ToolParameter.ToolParameterForm.FORM, + llm_description="Sampling method of Stable Diffusion, you can check the official documentation" + " of Stable Diffusion", + required=True, + default=sample_methods[0], + options=[ + ToolParameterOption(value=i, label=I18nObject(en_US=i, zh_Hans=i)) for i in sample_methods + ], + ) + ) + return parameters diff --git a/api/core/tools/provider/builtin/stablediffusion/tools/stable_diffusion.yaml b/api/core/tools/provider/builtin/stablediffusion/tools/stable_diffusion.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bbbdb16caf21bb13262b29bca05a8c063bf3d102 --- /dev/null +++ b/api/core/tools/provider/builtin/stablediffusion/tools/stable_diffusion.yaml @@ -0,0 +1,104 @@ +identity: + name: stable_diffusion + author: Dify + label: + en_US: Stable Diffusion WebUI + zh_Hans: Stable Diffusion WebUI + pt_BR: Stable Diffusion WebUI +description: + human: + en_US: A tool for generating images which can be deployed locally, you can use stable-diffusion-webui to deploy it. + zh_Hans: 一个可以在本地部署的图片生成的工具,您可以使用 stable-diffusion-webui 来部署它。 + pt_BR: A tool for generating images which can be deployed locally, you can use stable-diffusion-webui to deploy it. + llm: draw the image you want based on your prompt. +parameters: + - name: prompt + type: string + required: true + label: + en_US: Prompt + zh_Hans: 提示词 + pt_BR: Prompt + human_description: + en_US: Image prompt, you can check the official documentation of Stable Diffusion + zh_Hans: 图像提示词,您可以查看 Stable Diffusion 的官方文档 + pt_BR: Image prompt, you can check the official documentation of Stable Diffusion + llm_description: Image prompt of Stable Diffusion, you should describe the image you want to generate as a list of words as possible as detailed, the prompt must be written in English. + form: llm + - name: model + type: string + required: false + label: + en_US: Model Name + zh_Hans: 模型名称 + pt_BR: Model Name + human_description: + en_US: Model Name + zh_Hans: 模型名称 + pt_BR: Model Name + form: form + - name: lora + type: string + required: false + label: + en_US: Lora + zh_Hans: Lora + pt_BR: Lora + human_description: + en_US: Lora + zh_Hans: Lora + pt_BR: Lora + form: form + default: "" + - name: steps + type: number + required: false + label: + en_US: Steps + zh_Hans: Steps + pt_BR: Steps + human_description: + en_US: Steps + zh_Hans: Steps + pt_BR: Steps + form: form + default: 10 + - name: width + type: number + required: false + label: + en_US: Width + zh_Hans: Width + pt_BR: Width + human_description: + en_US: Width + zh_Hans: Width + pt_BR: Width + form: form + default: 1024 + - name: height + type: number + required: false + label: + en_US: Height + zh_Hans: Height + pt_BR: Height + human_description: + en_US: Height + zh_Hans: Height + pt_BR: Height + form: form + default: 1024 + - name: negative_prompt + type: string + required: false + label: + en_US: Negative prompt + zh_Hans: Negative prompt + pt_BR: Negative prompt + human_description: + en_US: Negative prompt + zh_Hans: Negative prompt + pt_BR: Negative prompt + form: form + default: bad art, ugly, deformed, watermark, duplicated, discontinuous lines diff --git a/api/core/tools/provider/builtin/stackexchange/_assets/icon.svg b/api/core/tools/provider/builtin/stackexchange/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..7042bc0e4156c948913fa560c173fbaaf41af6d5 --- /dev/null +++ b/api/core/tools/provider/builtin/stackexchange/_assets/icon.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/stackexchange/stackexchange.py b/api/core/tools/provider/builtin/stackexchange/stackexchange.py new file mode 100644 index 0000000000000000000000000000000000000000..9680c633cc701c9124532a63356f411952f5f747 --- /dev/null +++ b/api/core/tools/provider/builtin/stackexchange/stackexchange.py @@ -0,0 +1,25 @@ +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.stackexchange.tools.searchStackExQuestions import SearchStackExQuestionsTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class StackExchangeProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + try: + SearchStackExQuestionsTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "intitle": "Test", + "sort": "relevance", + "order": "desc", + "site": "stackoverflow", + "accepted": True, + "pagesize": 1, + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/stackexchange/stackexchange.yaml b/api/core/tools/provider/builtin/stackexchange/stackexchange.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d382a3cca9cef29cf7c4a13c7f0a8b200d6f799e --- /dev/null +++ b/api/core/tools/provider/builtin/stackexchange/stackexchange.yaml @@ -0,0 +1,13 @@ +identity: + author: Richards Tu + name: stackexchange + label: + en_US: Stack Exchange + zh_Hans: Stack Exchange + description: + en_US: Access questions and answers from the Stack Exchange and its sub-sites. + zh_Hans: 从 Stack Exchange 和其子论坛获取问题和答案。 + icon: icon.svg + tags: + - search + - utilities diff --git a/api/core/tools/provider/builtin/stackexchange/tools/fetchAnsByStackExQuesID.py b/api/core/tools/provider/builtin/stackexchange/tools/fetchAnsByStackExQuesID.py new file mode 100644 index 0000000000000000000000000000000000000000..534532009501f5ba3ff44b5dbc37a9fd9c6bfc34 --- /dev/null +++ b/api/core/tools/provider/builtin/stackexchange/tools/fetchAnsByStackExQuesID.py @@ -0,0 +1,39 @@ +from typing import Any, Union + +import requests +from pydantic import BaseModel, Field + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class FetchAnsByStackExQuesIDInput(BaseModel): + id: int = Field(..., description="The question ID") + site: str = Field(..., description="The Stack Exchange site") + order: str = Field(..., description="asc or desc") + sort: str = Field(..., description="activity, votes, creation") + pagesize: int = Field(..., description="Number of answers per page") + page: int = Field(..., description="Page number") + + +class FetchAnsByStackExQuesIDTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + input = FetchAnsByStackExQuesIDInput(**tool_parameters) + + params = { + "site": input.site, + "filter": "!nNPvSNdWme", + "order": input.order, + "sort": input.sort, + "pagesize": input.pagesize, + "page": input.page, + } + + response = requests.get(f"https://api.stackexchange.com/2.3/questions/{input.id}/answers", params=params) + + if response.status_code == 200: + return self.create_text_message(self.summary(user_id=user_id, content=response.text)) + else: + return self.create_text_message(f"API request failed with status code {response.status_code}") diff --git a/api/core/tools/provider/builtin/stackexchange/tools/fetchAnsByStackExQuesID.yaml b/api/core/tools/provider/builtin/stackexchange/tools/fetchAnsByStackExQuesID.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d663bce6097441f42d5e380a7334378f6107dbae --- /dev/null +++ b/api/core/tools/provider/builtin/stackexchange/tools/fetchAnsByStackExQuesID.yaml @@ -0,0 +1,107 @@ +identity: + name: fetchAnsByStackExQuesID + author: Richards Tu + label: + en_US: Fetch Stack Exchange Answers + zh_Hans: 获取 Stack Exchange 答案 +description: + human: + en_US: A tool for retrieving answers for a specific Stack Exchange question ID. Must be used with the searchStackExQuesID tool. + zh_Hans: 用于检索特定Stack Exchange问题ID的答案的工具。必须与searchStackExQuesID工具一起使用。 + llm: A tool for retrieving answers for Stack Exchange question ID. +parameters: + - name: id + type: string + required: true + label: + en_US: Question ID + zh_Hans: 问题ID + human_description: + en_US: The ID of the Stack Exchange question to fetch answers for. + zh_Hans: 要获取答案的Stack Exchange问题的ID。 + llm_description: The ID of the Stack Exchange question. + form: llm + - name: site + type: string + required: true + label: + en_US: Stack Exchange site + zh_Hans: Stack Exchange站点 + human_description: + en_US: The Stack Exchange site the question is from, e.g. stackoverflow, unix, etc. + zh_Hans: 问题所在的Stack Exchange站点,例如stackoverflow、unix等。 + llm_description: Stack Exchange site identifier - 'stackoverflow', 'serverfault', 'superuser', 'askubuntu', 'unix', 'cs', 'softwareengineering', 'codegolf', 'codereview', 'cstheory', 'security', 'cryptography', 'reverseengineering', 'datascience', 'devops', 'ux', 'dba', 'gis', 'webmasters', 'arduino', 'raspberrypi', 'networkengineering', 'iot', 'tor', 'sqa', 'mathoverflow', 'math', 'mathematica', 'dsp', 'gamedev', 'robotics', 'genai', 'computergraphics'. + form: llm + - name: filter + type: string + required: true + label: + en_US: Filter + zh_Hans: 过滤器 + human_description: + en_US: This is required in order to actually get the body of the answer. + zh_Hans: 为了实际获取答案的正文是必需的。 + options: + - value: "!nNPvSNdWme" + label: + en_US: Must Select + zh_Hans: 必须选择 + form: form + default: "!nNPvSNdWme" + - name: order + type: string + required: true + label: + en_US: Sort direction + zh_Hans: 排序方向 + human_description: + en_US: The direction to sort the answers - ascending or descending. + zh_Hans: 答案的排序方向 - 升序或降序。 + form: form + options: + - value: asc + label: + en_US: Ascending + zh_Hans: 升序 + - value: desc + label: + en_US: Descending + zh_Hans: 降序 + default: desc + - name: sort + type: string + required: true + label: + en_US: Sort order + zh_Hans: 排序 + human_description: + en_US: The sort order for the answers - activity, votes, or creation date. + zh_Hans: 答案的排序顺序 - 活动、投票或创建日期。 + llm_description: activity, votes, or creation. + form: llm + - name: pagesize + type: number + required: true + label: + en_US: Results per page + zh_Hans: 每页结果数 + human_description: + en_US: The number of answers to return per page. + zh_Hans: 每页返回的答案数。 + form: form + min: 1 + max: 5 + default: 1 + - name: page + type: number + required: true + label: + en_US: Page number + zh_Hans: 页码 + human_description: + en_US: The page number of answers to retrieve. + zh_Hans: 要检索的答案的页码。 + form: form + min: 1 + max: 5 + default: 3 diff --git a/api/core/tools/provider/builtin/stackexchange/tools/searchStackExQuestions.py b/api/core/tools/provider/builtin/stackexchange/tools/searchStackExQuestions.py new file mode 100644 index 0000000000000000000000000000000000000000..4a25a808adf26a00699c0fc99f4ae8b551ee6091 --- /dev/null +++ b/api/core/tools/provider/builtin/stackexchange/tools/searchStackExQuestions.py @@ -0,0 +1,45 @@ +from typing import Any, Union + +import requests +from pydantic import BaseModel, Field + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class SearchStackExQuestionsInput(BaseModel): + intitle: str = Field(..., description="The search query.") + sort: str = Field(..., description="The sort order - relevance, activity, votes, creation.") + order: str = Field(..., description="asc or desc") + site: str = Field(..., description="The Stack Exchange site.") + tagged: str = Field(None, description="Semicolon-separated tags to include.") + nottagged: str = Field(None, description="Semicolon-separated tags to exclude.") + accepted: bool = Field(..., description="true for only accepted answers, false otherwise") + pagesize: int = Field(..., description="Number of results per page") + + +class SearchStackExQuestionsTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + input = SearchStackExQuestionsInput(**tool_parameters) + + params = { + "intitle": input.intitle, + "sort": input.sort, + "order": input.order, + "site": input.site, + "accepted": input.accepted, + "pagesize": input.pagesize, + } + if input.tagged: + params["tagged"] = input.tagged + if input.nottagged: + params["nottagged"] = input.nottagged + + response = requests.get("https://api.stackexchange.com/2.3/search", params=params) + + if response.status_code == 200: + return self.create_text_message(self.summary(user_id=user_id, content=response.text)) + else: + return self.create_text_message(f"API request failed with status code {response.status_code}") diff --git a/api/core/tools/provider/builtin/stackexchange/tools/searchStackExQuestions.yaml b/api/core/tools/provider/builtin/stackexchange/tools/searchStackExQuestions.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bbfbae38b06e1a893bbf56780a7abccf47753fcb --- /dev/null +++ b/api/core/tools/provider/builtin/stackexchange/tools/searchStackExQuestions.yaml @@ -0,0 +1,121 @@ +identity: + name: searchStackExQuestions + author: Richards Tu + label: + en_US: Search Stack Exchange Questions + zh_Hans: 搜索Stack Exchange问题 +description: + human: + en_US: A tool for searching questions on a Stack Exchange site. + zh_Hans: 在Stack Exchange站点上搜索问题的工具。 + llm: A tool for searching questions on Stack Exchange site. +parameters: + - name: intitle + type: string + required: true + label: + en_US: Search query + zh_Hans: 搜索查询 + human_description: + en_US: The search query to use for finding questions. + zh_Hans: 用于查找问题的搜索查询。 + llm_description: The search query. + form: llm + - name: sort + type: string + required: true + label: + en_US: Sort order + zh_Hans: 排序 + human_description: + en_US: The sort order for the search results - relevance, activity, votes, or creation date. + zh_Hans: 搜索结果的排序顺序 - 相关性、活动、投票或创建日期。 + llm_description: The sort order - 'relevance', 'activity', 'votes', or 'creation'. + form: llm + - name: order + type: select + required: true + label: + en_US: Sort direction + zh_Hans: 排序方向 + human_description: + en_US: The direction to sort - ascending or descending. + zh_Hans: 排序方向 - 升序或降序。 + form: form + options: + - value: asc + label: + en_US: Ascending + zh_Hans: 升序 + - value: desc + label: + en_US: Descending + zh_Hans: 降序 + default: desc + - name: site + type: string + required: true + label: + en_US: Stack Exchange site + zh_Hans: Stack Exchange 站点 + human_description: + en_US: The Stack Exchange site to search, e.g. stackoverflow, unix, etc. + zh_Hans: 要搜索的Stack Exchange站点,例如stackoverflow、unix等。 + llm_description: Stack Exchange site identifier - 'stackoverflow', 'serverfault', 'superuser', 'askubuntu', 'unix', 'cs', 'softwareengineering', 'codegolf', 'codereview', 'cstheory', 'security', 'cryptography', 'reverseengineering', 'datascience', 'devops', 'ux', 'dba', 'gis', 'webmasters', 'arduino', 'raspberrypi', 'networkengineering', 'iot', 'tor', 'sqa', 'mathoverflow', 'math', 'mathematica', 'dsp', 'gamedev', 'robotics', 'genai', 'computergraphics'. + form: llm + - name: tagged + type: string + required: false + label: + en_US: Include tags + zh_Hans: 包含标签 + human_description: + en_US: A semicolon-separated list of tags that questions must have. + zh_Hans: 问题必须具有的标签的分号分隔列表。 + llm_description: Semicolon-separated tags to include. Leave blank if not needed. + form: llm + - name: nottagged + type: string + required: false + label: + en_US: Exclude tags + zh_Hans: 排除标签 + human_description: + en_US: A semicolon-separated list of tags to exclude from the search. + zh_Hans: 从搜索中排除的标签的分号分隔列表。 + llm_description: Semicolon-separated tags to exclude. Leave blank if not needed. + form: llm + - name: accepted + type: boolean + required: true + label: + en_US: Has accepted answer + zh_Hans: 有已接受的答案 + human_description: + en_US: Whether to limit to only questions that have an accepted answer. + zh_Hans: 是否限制为只有已接受答案的问题。 + form: form + options: + - value: 'true' + label: + en_US: 'Yes' + zh_Hans: 是 + - value: 'false' + label: + en_US: 'No' + zh_Hans: 否 + default: 'true' + - name: pagesize + type: number + required: true + label: + en_US: Results per page + zh_Hans: 每页结果数 + human_description: + en_US: The number of results to return per page. + zh_Hans: 每页返回的结果数。 + llm_description: The number of results per page. + form: form + min: 1 + max: 50 + default: 10 diff --git a/api/core/tools/provider/builtin/stepfun/__init__.py b/api/core/tools/provider/builtin/stepfun/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/tools/provider/builtin/stepfun/_assets/icon.png b/api/core/tools/provider/builtin/stepfun/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..85b96d0c74c24c2c28ccd0d363f02b35e359f561 Binary files /dev/null and b/api/core/tools/provider/builtin/stepfun/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/stepfun/stepfun.py b/api/core/tools/provider/builtin/stepfun/stepfun.py new file mode 100644 index 0000000000000000000000000000000000000000..239db85b1118b02c6141e12c82e0409f353a0ea5 --- /dev/null +++ b/api/core/tools/provider/builtin/stepfun/stepfun.py @@ -0,0 +1,24 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.stepfun.tools.image import StepfunTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class StepfunProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + StepfunTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "prompt": "cute girl, blue eyes, white hair, anime style", + "size": "256x256", + "n": 1, + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/stepfun/stepfun.yaml b/api/core/tools/provider/builtin/stepfun/stepfun.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e8139a4d7d6cfda340978c488ecae2d13a59937e --- /dev/null +++ b/api/core/tools/provider/builtin/stepfun/stepfun.yaml @@ -0,0 +1,33 @@ +identity: + author: Stepfun + name: stepfun + label: + en_US: Image-1X + zh_Hans: 阶跃星辰绘画 + description: + en_US: Image-1X + zh_Hans: 阶跃星辰绘画 + icon: icon.png + tags: + - image + - productivity +credentials_for_provider: + stepfun_api_key: + type: secret-input + required: true + label: + en_US: Stepfun API key + zh_Hans: 阶跃星辰API key + placeholder: + en_US: Please input your Stepfun API key + zh_Hans: 请输入你的阶跃星辰 API key + url: https://platform.stepfun.com/interface-key + stepfun_base_url: + type: text-input + required: false + label: + en_US: Stepfun base URL + zh_Hans: 阶跃星辰 base URL + placeholder: + en_US: Please input your Stepfun base URL + zh_Hans: 请输入你的阶跃星辰 base URL diff --git a/api/core/tools/provider/builtin/stepfun/tools/image.py b/api/core/tools/provider/builtin/stepfun/tools/image.py new file mode 100644 index 0000000000000000000000000000000000000000..61cc14fac6ca93275e8295ca9980e918be26ad45 --- /dev/null +++ b/api/core/tools/provider/builtin/stepfun/tools/image.py @@ -0,0 +1,66 @@ +from typing import Any, Union + +from openai import OpenAI +from yarl import URL + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class StepfunTool(BuiltinTool): + """Stepfun Image Generation Tool""" + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + base_url = self.runtime.credentials.get("stepfun_base_url") or "https://api.stepfun.com" + base_url = str(URL(base_url) / "v1") + + client = OpenAI( + api_key=self.runtime.credentials["stepfun_api_key"], + base_url=base_url, + ) + + extra_body = {} + model = "step-1x-medium" + # prompt + prompt = tool_parameters.get("prompt", "") + if not prompt: + return self.create_text_message("Please input prompt") + if len(prompt) > 1024: + return self.create_text_message("The prompt length should less than 1024") + seed = tool_parameters.get("seed", 0) + if seed > 0: + extra_body["seed"] = seed + steps = tool_parameters.get("steps", 50) + if steps > 0: + extra_body["steps"] = steps + cfg_scale = tool_parameters.get("cfg_scale", 7.5) + if cfg_scale > 0: + extra_body["cfg_scale"] = cfg_scale + + # call openapi stepfun model + response = client.images.generate( + prompt=prompt, + model=model, + size=tool_parameters.get("size", "1024x1024"), + n=tool_parameters.get("n", 1), + extra_body=extra_body, + ) + + result = [] + for image in response.data: + result.append(self.create_image_message(image=image.url)) + result.append( + self.create_json_message( + { + "url": image.url, + } + ) + ) + return result diff --git a/api/core/tools/provider/builtin/stepfun/tools/image.yaml b/api/core/tools/provider/builtin/stepfun/tools/image.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dfda6ed1914848ecca8ea6eb5840b8cd8955353f --- /dev/null +++ b/api/core/tools/provider/builtin/stepfun/tools/image.yaml @@ -0,0 +1,133 @@ +identity: + name: stepfun + author: Stepfun + label: + en_US: step-1x + zh_Hans: 阶跃星辰绘画 + pt_BR: step-1x + description: + en_US: step-1x is a powerful drawing tool by stepfun, you can draw the image based on your prompt + zh_Hans: step-1x 系列是阶跃星辰提供的强大的绘画工具,它可以根据您的提示词绘制出您想要的图像。 + pt_BR: step-1x is a powerful drawing tool by stepfun, you can draw the image based on your prompt +description: + human: + en_US: step-1x is a text to image tool + zh_Hans: step-1x 是一个文本/图像到图像的工具 + pt_BR: step-1x is a text to image tool + llm: step-1x is a tool used to generate images from text or image +parameters: + - name: prompt + type: string + required: true + label: + en_US: Prompt + zh_Hans: 提示词 + pt_BR: Prompt + human_description: + en_US: Image prompt, you can check the official documentation of step-1x + zh_Hans: 图像提示词,您可以查看 step-1x 的官方文档 + pt_BR: Image prompt, you can check the official documentation of step-1x + llm_description: Image prompt of step-1x you should describe the image you want to generate as a list of words as possible as detailed + form: llm + - name: size + type: select + required: false + human_description: + en_US: The size of the generated image + zh_Hans: 生成的图片大小 + pt_BR: The size of the generated image + label: + en_US: Image size + zh_Hans: 图像大小 + pt_BR: Image size + form: form + options: + - value: 256x256 + label: + en_US: 256x256 + zh_Hans: 256x256 + pt_BR: 256x256 + - value: 512x512 + label: + en_US: 512x512 + zh_Hans: 512x512 + pt_BR: 512x512 + - value: 768x768 + label: + en_US: 768x768 + zh_Hans: 768x768 + pt_BR: 768x768 + - value: 1024x1024 + label: + en_US: 1024x1024 + zh_Hans: 1024x1024 + pt_BR: 1024x1024 + - value: 1280x800 + label: + en_US: 1280x800 + zh_Hans: 1280x800 + pt_BR: 1280x800 + - value: 800x1280 + label: + en_US: 800x1280 + zh_Hans: 800x1280 + pt_BR: 800x1280 + default: 1024x1024 + - name: n + type: number + required: true + human_description: + en_US: Number of generated images, now only one image can be generated at a time + zh_Hans: 生成的图像数量,当前仅支持每次生成一张图片 + pt_BR: Number of generated images, now only one image can be generated at a time + label: + en_US: Number of generated images + zh_Hans: 生成的图像数量 + pt_BR: Number of generated images + form: form + default: 1 + min: 1 + max: 1 + - name: seed + type: number + required: false + label: + en_US: seed + zh_Hans: seed + pt_BR: seed + human_description: + en_US: seed + zh_Hans: seed + pt_BR: seed + form: form + default: 10 + - name: steps + type: number + required: false + label: + en_US: Steps + zh_Hans: Steps + pt_BR: Steps + human_description: + en_US: Steps, now support integers between 1 and 100 + zh_Hans: Steps, 当前支持 1~100 之间整数 + pt_BR: Steps, now support integers between 1 and 100 + form: form + default: 50 + min: 1 + max: 100 + - name: cfg_scale + type: number + required: false + label: + en_US: classifier-free guidance scale + zh_Hans: classifier-free guidance scale + pt_BR: classifier-free guidance scale + human_description: + en_US: classifier-free guidance scale + zh_Hans: classifier-free guidance scale + pt_BR: classifier-free guidance scale + form: form + default: 7.5 + min: 1 + max: 10 diff --git a/api/core/tools/provider/builtin/tavily/_assets/icon.png b/api/core/tools/provider/builtin/tavily/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..fdb40ab5689ba9f40b22d2c700ed2ce1b2602829 Binary files /dev/null and b/api/core/tools/provider/builtin/tavily/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/tavily/tavily.py b/api/core/tools/provider/builtin/tavily/tavily.py new file mode 100644 index 0000000000000000000000000000000000000000..a702b0a74e6131694c479df784a4408a152d21e0 --- /dev/null +++ b/api/core/tools/provider/builtin/tavily/tavily.py @@ -0,0 +1,29 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.tavily.tools.tavily_search import TavilySearchTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class TavilyProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + TavilySearchTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "query": "Sachin Tendulkar", + "search_depth": "basic", + "include_answer": True, + "include_images": False, + "include_raw_content": False, + "max_results": 5, + "include_domains": "", + "exclude_domains": "", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/tavily/tavily.yaml b/api/core/tools/provider/builtin/tavily/tavily.yaml new file mode 100644 index 0000000000000000000000000000000000000000..aba621b094e81665449aeb252f57e0eb7ddaa9a3 --- /dev/null +++ b/api/core/tools/provider/builtin/tavily/tavily.yaml @@ -0,0 +1,26 @@ +identity: + author: Yash Parmar, Kalo Chin + name: tavily + label: + en_US: Tavily Search & Extract + zh_Hans: Tavily 搜索和提取 + description: + en_US: A powerful AI-native search engine and web content extraction tool that provides highly relevant search results and raw content extraction from web pages. + zh_Hans: 一个强大的原生AI搜索引擎和网页内容提取工具,提供高度相关的搜索结果和网页原始内容提取。 + icon: icon.png + tags: + - search +credentials_for_provider: + tavily_api_key: + type: secret-input + required: true + label: + en_US: Tavily API key + zh_Hans: Tavily API key + placeholder: + en_US: Please input your Tavily API key + zh_Hans: 请输入你的 Tavily API key + help: + en_US: Get your Tavily API key from Tavily + zh_Hans: 从 TavilyApi 获取您的 Tavily API key + url: https://app.tavily.com/home diff --git a/api/core/tools/provider/builtin/tavily/tools/tavily_extract.py b/api/core/tools/provider/builtin/tavily/tools/tavily_extract.py new file mode 100644 index 0000000000000000000000000000000000000000..a37548018d44ff39ce08190a12bb211e2c59b349 --- /dev/null +++ b/api/core/tools/provider/builtin/tavily/tools/tavily_extract.py @@ -0,0 +1,145 @@ +from typing import Any + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +TAVILY_API_URL = "https://api.tavily.com" + + +class TavilyExtract: + """ + A class for extracting content from web pages using the Tavily Extract API. + + Args: + api_key (str): The API key for accessing the Tavily Extract API. + + Methods: + extract_content: Retrieves extracted content from the Tavily Extract API. + """ + + def __init__(self, api_key: str) -> None: + self.api_key = api_key + + def extract_content(self, params: dict[str, Any]) -> dict: + """ + Retrieves extracted content from the Tavily Extract API. + + Args: + params (Dict[str, Any]): The extraction parameters. + + Returns: + dict: The extracted content. + + """ + # Ensure required parameters are set + if "api_key" not in params: + params["api_key"] = self.api_key + + # Process parameters + processed_params = self._process_params(params) + + response = requests.post(f"{TAVILY_API_URL}/extract", json=processed_params) + response.raise_for_status() + return response.json() + + def _process_params(self, params: dict[str, Any]) -> dict: + """ + Processes and validates the extraction parameters. + + Args: + params (Dict[str, Any]): The extraction parameters. + + Returns: + dict: The processed parameters. + """ + processed_params = {} + + # Process 'urls' + if "urls" in params: + urls = params["urls"] + if isinstance(urls, str): + processed_params["urls"] = [url.strip() for url in urls.replace(",", " ").split()] + elif isinstance(urls, list): + processed_params["urls"] = urls + else: + raise ValueError("The 'urls' parameter is required.") + + # Only include 'api_key' + processed_params["api_key"] = params.get("api_key", self.api_key) + + return processed_params + + +class TavilyExtractTool(BuiltinTool): + """ + A tool for extracting content from web pages using Tavily Extract. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage | list[ToolInvokeMessage]: + """ + Invokes the Tavily Extract tool with the given user ID and tool parameters. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (Dict[str, Any]): The parameters for the Tavily Extract tool. + + Returns: + ToolInvokeMessage | list[ToolInvokeMessage]: The result of the Tavily Extract tool invocation. + """ + urls = tool_parameters.get("urls", "") + api_key = self.runtime.credentials.get("tavily_api_key") + if not api_key: + return self.create_text_message( + "Tavily API key is missing. Please set the 'tavily_api_key' in credentials." + ) + if not urls: + return self.create_text_message("Please input at least one URL to extract.") + + tavily_extract = TavilyExtract(api_key) + try: + raw_results = tavily_extract.extract_content(tool_parameters) + except requests.HTTPError as e: + return self.create_text_message(f"Error occurred while extracting content: {str(e)}") + + if not raw_results.get("results"): + return self.create_text_message("No content could be extracted from the provided URLs.") + else: + # Always return JSON message with all data + json_message = self.create_json_message(raw_results) + + # Create text message based on user-selected parameters + text_message_content = self._format_results_as_text(raw_results) + text_message = self.create_text_message(text=text_message_content) + + return [json_message, text_message] + + def _format_results_as_text(self, raw_results: dict) -> str: + """ + Formats the raw extraction results into a markdown text based on user-selected parameters. + + Args: + raw_results (dict): The raw extraction results. + + Returns: + str: The formatted markdown text. + """ + output_lines = [] + + for idx, result in enumerate(raw_results.get("results", []), 1): + url = result.get("url", "") + raw_content = result.get("raw_content", "") + + output_lines.append(f"## Extracted Content {idx}: {url}\n") + output_lines.append(f"**Raw Content:**\n{raw_content}\n") + output_lines.append("---\n") + + if raw_results.get("failed_results"): + output_lines.append("## Failed URLs:\n") + for failed in raw_results["failed_results"]: + url = failed.get("url", "") + error = failed.get("error", "Unknown error") + output_lines.append(f"- {url}: {error}\n") + + return "\n".join(output_lines) diff --git a/api/core/tools/provider/builtin/tavily/tools/tavily_extract.yaml b/api/core/tools/provider/builtin/tavily/tools/tavily_extract.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a04da73b540f4d747dee92f0f399906f32e3aa4c --- /dev/null +++ b/api/core/tools/provider/builtin/tavily/tools/tavily_extract.yaml @@ -0,0 +1,23 @@ +identity: + name: tavily_extract + author: Kalo Chin + label: + en_US: Tavily Extract + zh_Hans: Tavily Extract +description: + human: + en_US: A web extraction tool built specifically for AI agents (LLMs), delivering raw content from web pages. + zh_Hans: 专为人工智能代理 (LLM) 构建的网页提取工具,提供网页的原始内容。 + llm: A tool for extracting raw content from web pages, designed for AI agents (LLMs). +parameters: + - name: urls + type: string + required: true + label: + en_US: URLs + zh_Hans: URLs + human_description: + en_US: A comma-separated list of URLs to extract content from. + zh_Hans: 要从中提取内容的 URL 的逗号分隔列表。 + llm_description: A comma-separated list of URLs to extract content from. + form: llm diff --git a/api/core/tools/provider/builtin/tavily/tools/tavily_search.py b/api/core/tools/provider/builtin/tavily/tools/tavily_search.py new file mode 100644 index 0000000000000000000000000000000000000000..ea41ea3ca3c61f4676fb76860eeb5039ae6cd3a1 --- /dev/null +++ b/api/core/tools/provider/builtin/tavily/tools/tavily_search.py @@ -0,0 +1,195 @@ +from typing import Any + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +TAVILY_API_URL = "https://api.tavily.com" + + +class TavilySearch: + """ + A class for performing search operations using the Tavily Search API. + + Args: + api_key (str): The API key for accessing the Tavily Search API. + + Methods: + raw_results: Retrieves raw search results from the Tavily Search API. + """ + + def __init__(self, api_key: str) -> None: + self.api_key = api_key + + def raw_results(self, params: dict[str, Any]) -> dict: + """ + Retrieves raw search results from the Tavily Search API. + + Args: + params (Dict[str, Any]): The search parameters. + + Returns: + dict: The raw search results. + + """ + # Ensure required parameters are set + params["api_key"] = self.api_key + + # Process parameters to ensure correct types + processed_params = self._process_params(params) + + response = requests.post(f"{TAVILY_API_URL}/search", json=processed_params) + response.raise_for_status() + return response.json() + + def _process_params(self, params: dict[str, Any]) -> dict: + """ + Processes and validates the search parameters. + + Args: + params (Dict[str, Any]): The search parameters. + + Returns: + dict: The processed parameters. + """ + processed_params = {} + + for key, value in params.items(): + if value is None or value == "None": + continue + if key in ["include_domains", "exclude_domains"]: + if isinstance(value, str): + # Split the string by commas or spaces and strip whitespace + processed_params[key] = [domain.strip() for domain in value.replace(",", " ").split()] + elif key in ["include_images", "include_image_descriptions", "include_answer", "include_raw_content"]: + # Ensure boolean type + if isinstance(value, str): + processed_params[key] = value.lower() == "true" + else: + processed_params[key] = bool(value) + elif key in ["max_results", "days"]: + if isinstance(value, str): + processed_params[key] = int(value) + else: + processed_params[key] = value + elif key in ["search_depth", "topic", "query", "api_key"]: + processed_params[key] = value + else: + # Unrecognized parameter + pass + + # Set defaults if not present + processed_params.setdefault("search_depth", "basic") + processed_params.setdefault("topic", "general") + processed_params.setdefault("max_results", 5) + + # If topic is 'news', ensure 'days' is set + if processed_params.get("topic") == "news": + processed_params.setdefault("days", 3) + + return processed_params + + +class TavilySearchTool(BuiltinTool): + """ + A tool for searching Tavily using a given query. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage | list[ToolInvokeMessage]: + """ + Invokes the Tavily search tool with the given user ID and tool parameters. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (Dict[str, Any]): The parameters for the Tavily search tool. + + Returns: + ToolInvokeMessage | list[ToolInvokeMessage]: The result of the Tavily search tool invocation. + """ + query = tool_parameters.get("query", "") + api_key = self.runtime.credentials.get("tavily_api_key") + if not api_key: + return self.create_text_message( + "Tavily API key is missing. Please set the 'tavily_api_key' in credentials." + ) + if not query: + return self.create_text_message("Please input a query.") + + tavily_search = TavilySearch(api_key) + try: + raw_results = tavily_search.raw_results(tool_parameters) + except requests.HTTPError as e: + return self.create_text_message(f"Error occurred while searching: {str(e)}") + + if not raw_results.get("results"): + return self.create_text_message(f"No results found for '{query}' in Tavily.") + else: + # Always return JSON message with all data + json_message = self.create_json_message(raw_results) + + # Create text message based on user-selected parameters + text_message_content = self._format_results_as_text(raw_results, tool_parameters) + text_message = self.create_text_message(text=text_message_content) + + return [json_message, text_message] + + def _format_results_as_text(self, raw_results: dict, tool_parameters: dict[str, Any]) -> str: + """ + Formats the raw results into a markdown text based on user-selected parameters. + + Args: + raw_results (dict): The raw search results. + tool_parameters (dict): The tool parameters selected by the user. + + Returns: + str: The formatted markdown text. + """ + output_lines = [] + + # Include answer if requested + if tool_parameters.get("include_answer", False) and raw_results.get("answer"): + output_lines.append(f"**Answer:** {raw_results['answer']}\n") + + # Include images if requested + if tool_parameters.get("include_images", False) and raw_results.get("images"): + output_lines.append("**Images:**\n") + for image in raw_results["images"]: + if tool_parameters.get("include_image_descriptions", False) and "description" in image: + output_lines.append(f"![{image['description']}]({image['url']})\n") + else: + output_lines.append(f"![]({image['url']})\n") + + # Process each result + if "results" in raw_results: + for idx, result in enumerate(raw_results["results"], 1): + title = result.get("title", "No Title") + url = result.get("url", "") + content = result.get("content", "") + published_date = result.get("published_date", "") + score = result.get("score", "") + + output_lines.append(f"### Result {idx}: [{title}]({url})\n") + + # Include published date if available and topic is 'news' + if tool_parameters.get("topic") == "news" and published_date: + output_lines.append(f"**Published Date:** {published_date}\n") + + output_lines.append(f"**URL:** {url}\n") + + # Include score (relevance) + if score: + output_lines.append(f"**Relevance Score:** {score}\n") + + # Include content + if content: + output_lines.append(f"**Content:**\n{content}\n") + + # Include raw content if requested + if tool_parameters.get("include_raw_content", False) and result.get("raw_content"): + output_lines.append(f"**Raw Content:**\n{result['raw_content']}\n") + + # Add a separator + output_lines.append("---\n") + + return "\n".join(output_lines) diff --git a/api/core/tools/provider/builtin/tavily/tools/tavily_search.yaml b/api/core/tools/provider/builtin/tavily/tools/tavily_search.yaml new file mode 100644 index 0000000000000000000000000000000000000000..14b2829701fe48c11d30c476a7e101930a0581ef --- /dev/null +++ b/api/core/tools/provider/builtin/tavily/tools/tavily_search.yaml @@ -0,0 +1,152 @@ +identity: + name: tavily_search + author: Yash Parmar + label: + en_US: Tavily Search + zh_Hans: Tavily Search +description: + human: + en_US: A search engine tool built specifically for AI agents (LLMs), delivering real-time, accurate, and factual results at speed. + zh_Hans: 专为人工智能代理 (LLM) 构建的搜索引擎工具,可快速提供实时、准确和真实的结果。 + llm: A tool for search engine built specifically for AI agents (LLMs), delivering real-time, accurate, and factual results at speed. +parameters: + - name: query + type: string + required: true + label: + en_US: Query + zh_Hans: 查询 + human_description: + en_US: The search query you want to execute with Tavily. + zh_Hans: 您想用 Tavily 执行的搜索查询。 + llm_description: The search query. + form: llm + - name: search_depth + type: select + required: false + label: + en_US: Search Depth + zh_Hans: 搜索深度 + human_description: + en_US: The depth of the search. + zh_Hans: 搜索的深度。 + form: form + options: + - value: basic + label: + en_US: Basic + zh_Hans: 基本 + - value: advanced + label: + en_US: Advanced + zh_Hans: 高级 + default: basic + - name: topic + type: select + required: false + label: + en_US: Topic + zh_Hans: 主题 + human_description: + en_US: The category of the search. + zh_Hans: 搜索的类别。 + form: form + options: + - value: general + label: + en_US: General + zh_Hans: 一般 + - value: news + label: + en_US: News + zh_Hans: 新闻 + default: general + - name: days + type: number + required: false + label: + en_US: Days + zh_Hans: 天数 + human_description: + en_US: The number of days back from the current date to include in the search results (only applicable when "topic" is "news"). + zh_Hans: 从当前日期起向前追溯的天数,以包含在搜索结果中(仅当“topic”为“news”时适用)。 + form: form + min: 1 + default: 3 + - name: max_results + type: number + required: false + label: + en_US: Max Results + zh_Hans: 最大结果数 + human_description: + en_US: The maximum number of search results to return. + zh_Hans: 要返回的最大搜索结果数。 + form: form + min: 1 + max: 20 + default: 5 + - name: include_images + type: boolean + required: false + label: + en_US: Include Images + zh_Hans: 包含图片 + human_description: + en_US: Include a list of query-related images in the response. + zh_Hans: 在响应中包含与查询相关的图片列表。 + form: form + default: false + - name: include_image_descriptions + type: boolean + required: false + label: + en_US: Include Image Descriptions + zh_Hans: 包含图片描述 + human_description: + en_US: When include_images is True, adds descriptive text for each image. + zh_Hans: 当 include_images 为 True 时,为每个图像添加描述文本。 + form: form + default: false + - name: include_answer + type: boolean + required: false + label: + en_US: Include Answer + zh_Hans: 包含答案 + human_description: + en_US: Include a short answer to the original query in the response. + zh_Hans: 在响应中包含对原始查询的简短回答。 + form: form + default: false + - name: include_raw_content + type: boolean + required: false + label: + en_US: Include Raw Content + zh_Hans: 包含原始内容 + human_description: + en_US: Include the cleaned and parsed HTML content of each search result. + zh_Hans: 包含每个搜索结果的已清理和解析的HTML内容。 + form: form + default: false + - name: include_domains + type: string + required: false + label: + en_US: Include Domains + zh_Hans: 包含域 + human_description: + en_US: A comma-separated list of domains to specifically include in the search results. + zh_Hans: 要在搜索结果中特别包含的域的逗号分隔列表。 + form: form + - name: exclude_domains + type: string + required: false + label: + en_US: Exclude Domains + zh_Hans: 排除域 + human_description: + en_US: A comma-separated list of domains to specifically exclude from the search results. + zh_Hans: 要从搜索结果中特别排除的域的逗号分隔列表。 + form: form diff --git a/api/core/tools/provider/builtin/tianditu/_assets/icon.svg b/api/core/tools/provider/builtin/tianditu/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..749d4bda265ab02c204c5c28aa962a970bf30c6b --- /dev/null +++ b/api/core/tools/provider/builtin/tianditu/_assets/icon.svg @@ -0,0 +1,21 @@ + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/tianditu/tianditu.py b/api/core/tools/provider/builtin/tianditu/tianditu.py new file mode 100644 index 0000000000000000000000000000000000000000..cb7d7bd8bb2c412cbcd48cceaccc4d4724069fe0 --- /dev/null +++ b/api/core/tools/provider/builtin/tianditu/tianditu.py @@ -0,0 +1,23 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.tianditu.tools.poisearch import PoiSearchTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class TiandituProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + PoiSearchTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "content": "北京", + "specify": "156110000", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/tianditu/tianditu.yaml b/api/core/tools/provider/builtin/tianditu/tianditu.yaml new file mode 100644 index 0000000000000000000000000000000000000000..77af834bdc589385621e3b2a1701fba67340f740 --- /dev/null +++ b/api/core/tools/provider/builtin/tianditu/tianditu.yaml @@ -0,0 +1,32 @@ +identity: + author: Listeng + name: tianditu + label: + en_US: Tianditu + zh_Hans: 天地图 + pt_BR: Tianditu + description: + en_US: The Tianditu tool provided the functions of place name search, geocoding, static maps generation, etc. in China region. + zh_Hans: 天地图工具可以调用天地图的接口,实现中国区域内的地名搜索、地理编码、静态地图等功能。 + pt_BR: The Tianditu tool provided the functions of place name search, geocoding, static maps generation, etc. in China region. + icon: icon.svg + tags: + - utilities + - travel +credentials_for_provider: + tianditu_api_key: + type: secret-input + required: true + label: + en_US: Tianditu API Key + zh_Hans: 天地图Key + pt_BR: Tianditu API key + placeholder: + en_US: Please input your Tianditu API key + zh_Hans: 请输入你的天地图Key + pt_BR: Please input your Tianditu API key + help: + en_US: Get your Tianditu API key from Tianditu + zh_Hans: 获取您的天地图Key + pt_BR: Get your Tianditu API key from Tianditu + url: http://lbs.tianditu.gov.cn/home.html diff --git a/api/core/tools/provider/builtin/tianditu/tools/geocoder.py b/api/core/tools/provider/builtin/tianditu/tools/geocoder.py new file mode 100644 index 0000000000000000000000000000000000000000..690a0aed6f5affc6c8810266c8d2a8398411a725 --- /dev/null +++ b/api/core/tools/provider/builtin/tianditu/tools/geocoder.py @@ -0,0 +1,33 @@ +import json +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class GeocoderTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + base_url = "http://api.tianditu.gov.cn/geocoder" + + keyword = tool_parameters.get("keyword", "") + if not keyword: + return self.create_text_message("Invalid parameter keyword") + + tk = self.runtime.credentials["tianditu_api_key"] + + params = { + "keyWord": keyword, + } + + result = requests.get(base_url + "?ds=" + json.dumps(params, ensure_ascii=False) + "&tk=" + tk).json() + + return self.create_json_message(result) diff --git a/api/core/tools/provider/builtin/tianditu/tools/geocoder.yaml b/api/core/tools/provider/builtin/tianditu/tools/geocoder.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d6a168f9502019ddaf9fb5f35bd3989e1112b5db --- /dev/null +++ b/api/core/tools/provider/builtin/tianditu/tools/geocoder.yaml @@ -0,0 +1,26 @@ +identity: + name: geocoder + author: Listeng + label: + en_US: Get coords converted from address name + zh_Hans: 地理编码 + pt_BR: Get coords converted from address name +description: + human: + en_US: Geocoder + zh_Hans: 中国区域地理编码查询 + pt_BR: Geocoder + llm: A tool for geocoder in China +parameters: + - name: keyword + type: string + required: true + label: + en_US: keyword + zh_Hans: 搜索的关键字 + pt_BR: keyword + human_description: + en_US: keyword + zh_Hans: 搜索的关键字 + pt_BR: keyword + form: llm diff --git a/api/core/tools/provider/builtin/tianditu/tools/poisearch.py b/api/core/tools/provider/builtin/tianditu/tools/poisearch.py new file mode 100644 index 0000000000000000000000000000000000000000..798dd94d335654e6b949c8516db8216bfc3b35bc --- /dev/null +++ b/api/core/tools/provider/builtin/tianditu/tools/poisearch.py @@ -0,0 +1,58 @@ +import json +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class PoiSearchTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + geocoder_base_url = "http://api.tianditu.gov.cn/geocoder" + base_url = "http://api.tianditu.gov.cn/v2/search" + + keyword = tool_parameters.get("keyword", "") + if not keyword: + return self.create_text_message("Invalid parameter keyword") + + baseAddress = tool_parameters.get("baseAddress", "") + if not baseAddress: + return self.create_text_message("Invalid parameter baseAddress") + + tk = self.runtime.credentials["tianditu_api_key"] + + base_coords = requests.get( + geocoder_base_url + + "?ds=" + + json.dumps( + { + "keyWord": baseAddress, + }, + ensure_ascii=False, + ) + + "&tk=" + + tk + ).json() + + params = { + "keyWord": keyword, + "queryRadius": 5000, + "queryType": 3, + "pointLonlat": base_coords["location"]["lon"] + "," + base_coords["location"]["lat"], + "start": 0, + "count": 100, + } + + result = requests.get( + base_url + "?postStr=" + json.dumps(params, ensure_ascii=False) + "&type=query&tk=" + tk + ).json() + + return self.create_json_message(result) diff --git a/api/core/tools/provider/builtin/tianditu/tools/poisearch.yaml b/api/core/tools/provider/builtin/tianditu/tools/poisearch.yaml new file mode 100644 index 0000000000000000000000000000000000000000..01289d24e3d29a4f489f88d166f1271324852969 --- /dev/null +++ b/api/core/tools/provider/builtin/tianditu/tools/poisearch.yaml @@ -0,0 +1,38 @@ +identity: + name: point_of_interest_search + author: Listeng + label: + en_US: Point of Interest search + zh_Hans: 兴趣点搜索 + pt_BR: Point of Interest search +description: + human: + en_US: Search for certain types of points of interest around a location + zh_Hans: 搜索某个位置周边的5公里内某种类型的兴趣点 + pt_BR: Search for certain types of points of interest around a location + llm: A tool for searching for certain types of points of interest around a location +parameters: + - name: keyword + type: string + required: true + label: + en_US: poi keyword + zh_Hans: 兴趣点的关键字 + pt_BR: poi keyword + human_description: + en_US: poi keyword + zh_Hans: 兴趣点的关键字 + pt_BR: poi keyword + form: llm + - name: baseAddress + type: string + required: true + label: + en_US: base current point + zh_Hans: 当前位置的关键字 + pt_BR: base current point + human_description: + en_US: base current point + zh_Hans: 当前位置的关键字 + pt_BR: base current point + form: llm diff --git a/api/core/tools/provider/builtin/tianditu/tools/staticmap.py b/api/core/tools/provider/builtin/tianditu/tools/staticmap.py new file mode 100644 index 0000000000000000000000000000000000000000..aeaef08805768671bc3b8156c00711f41f0b6ac0 --- /dev/null +++ b/api/core/tools/provider/builtin/tianditu/tools/staticmap.py @@ -0,0 +1,49 @@ +import json +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class PoiSearchTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + + geocoder_base_url = "http://api.tianditu.gov.cn/geocoder" + base_url = "http://api.tianditu.gov.cn/staticimage" + + keyword = tool_parameters.get("keyword", "") + if not keyword: + return self.create_text_message("Invalid parameter keyword") + + tk = self.runtime.credentials["tianditu_api_key"] + + keyword_coords = requests.get( + geocoder_base_url + + "?ds=" + + json.dumps( + { + "keyWord": keyword, + }, + ensure_ascii=False, + ) + + "&tk=" + + tk + ).json() + coords = keyword_coords["location"]["lon"] + "," + keyword_coords["location"]["lat"] + + result = requests.get( + base_url + "?center=" + coords + "&markers=" + coords + "&width=400&height=300&zoom=14&tk=" + tk + ).content + + return self.create_blob_message( + blob=result, meta={"mime_type": "image/png"}, save_as=self.VariableKey.IMAGE.value + ) diff --git a/api/core/tools/provider/builtin/tianditu/tools/staticmap.yaml b/api/core/tools/provider/builtin/tianditu/tools/staticmap.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fc54c428066af51b3135b3fa4bdf55f1182e434c --- /dev/null +++ b/api/core/tools/provider/builtin/tianditu/tools/staticmap.yaml @@ -0,0 +1,26 @@ +identity: + name: generate_static_map + author: Listeng + label: + en_US: Generate a static map + zh_Hans: 生成静态地图 + pt_BR: Generate a static map +description: + human: + en_US: Generate a static map + zh_Hans: 生成静态地图 + pt_BR: Generate a static map + llm: A tool for generate a static map +parameters: + - name: keyword + type: string + required: true + label: + en_US: keyword + zh_Hans: 搜索的关键字 + pt_BR: keyword + human_description: + en_US: keyword + zh_Hans: 搜索的关键字 + pt_BR: keyword + form: llm diff --git a/api/core/tools/provider/builtin/time/_assets/icon.svg b/api/core/tools/provider/builtin/time/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..6d7118aed914ae9264316718e9609455f70c994a --- /dev/null +++ b/api/core/tools/provider/builtin/time/_assets/icon.svg @@ -0,0 +1,3 @@ + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/time/time.py b/api/core/tools/provider/builtin/time/time.py new file mode 100644 index 0000000000000000000000000000000000000000..e4df8d616cba381a0bbf705a902a88147f1ee6c2 --- /dev/null +++ b/api/core/tools/provider/builtin/time/time.py @@ -0,0 +1,16 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.time.tools.current_time import CurrentTimeTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class WikiPediaProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + CurrentTimeTool().invoke( + user_id="", + tool_parameters={}, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/time/time.yaml b/api/core/tools/provider/builtin/time/time.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1278939df589312a15e1475c267061962d5c5a6a --- /dev/null +++ b/api/core/tools/provider/builtin/time/time.yaml @@ -0,0 +1,15 @@ +identity: + author: Dify + name: time + label: + en_US: CurrentTime + zh_Hans: 时间 + pt_BR: CurrentTime + description: + en_US: A tool for getting the current time. + zh_Hans: 一个用于获取当前时间的工具。 + pt_BR: A tool for getting the current time. + icon: icon.svg + tags: + - utilities +credentials_for_provider: diff --git a/api/core/tools/provider/builtin/time/tools/current_time.py b/api/core/tools/provider/builtin/time/tools/current_time.py new file mode 100644 index 0000000000000000000000000000000000000000..6464bb6602b60a73c8d40404eb11ab7b3c4d9d6f --- /dev/null +++ b/api/core/tools/provider/builtin/time/tools/current_time.py @@ -0,0 +1,29 @@ +from datetime import UTC, datetime +from typing import Any, Union + +from pytz import timezone as pytz_timezone + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class CurrentTimeTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + # get timezone + tz = tool_parameters.get("timezone", "UTC") + fm = tool_parameters.get("format") or "%Y-%m-%d %H:%M:%S %Z" + if tz == "UTC": + return self.create_text_message(f"{datetime.now(UTC).strftime(fm)}") + + try: + tz = pytz_timezone(tz) + except: + return self.create_text_message(f"Invalid timezone: {tz}") + return self.create_text_message(f"{datetime.now(tz).strftime(fm)}") diff --git a/api/core/tools/provider/builtin/time/tools/current_time.yaml b/api/core/tools/provider/builtin/time/tools/current_time.yaml new file mode 100644 index 0000000000000000000000000000000000000000..52705ace4c1559571ee50f838668cd7b1ccf9dac --- /dev/null +++ b/api/core/tools/provider/builtin/time/tools/current_time.yaml @@ -0,0 +1,131 @@ +identity: + name: current_time + author: Dify + label: + en_US: Current Time + zh_Hans: 获取当前时间 + pt_BR: Current Time +description: + human: + en_US: A tool for getting the current time. + zh_Hans: 一个用于获取当前时间的工具。 + pt_BR: A tool for getting the current time. + llm: A tool for getting the current time. +parameters: + - name: format + type: string + required: false + label: + en_US: Format + zh_Hans: 格式 + pt_BR: Format + human_description: + en_US: Time format in strftime standard. + zh_Hans: strftime 标准的时间格式。 + pt_BR: Time format in strftime standard. + form: form + default: "%Y-%m-%d %H:%M:%S" + - name: timezone + type: select + required: false + label: + en_US: Timezone + zh_Hans: 时区 + pt_BR: Timezone + human_description: + en_US: Timezone + zh_Hans: 时区 + pt_BR: Timezone + form: form + default: UTC + options: + - value: UTC + label: + en_US: UTC + zh_Hans: UTC + pt_BR: UTC + - value: America/New_York + label: + en_US: America/New_York + zh_Hans: 美洲/纽约 + pt_BR: America/New_York + - value: America/Los_Angeles + label: + en_US: America/Los_Angeles + zh_Hans: 美洲/洛杉矶 + pt_BR: America/Los_Angeles + - value: America/Chicago + label: + en_US: America/Chicago + zh_Hans: 美洲/芝加哥 + pt_BR: America/Chicago + - value: America/Sao_Paulo + label: + en_US: America/Sao_Paulo + zh_Hans: 美洲/圣保罗 + pt_BR: América/São Paulo + - value: Asia/Shanghai + label: + en_US: Asia/Shanghai + zh_Hans: 亚洲/上海 + pt_BR: Asia/Shanghai + - value: Asia/Ho_Chi_Minh + label: + en_US: Asia/Ho_Chi_Minh + zh_Hans: 亚洲/胡志明市 + pt_BR: Ásia/Ho Chi Minh + - value: Asia/Tokyo + label: + en_US: Asia/Tokyo + zh_Hans: 亚洲/东京 + pt_BR: Asia/Tokyo + - value: Asia/Dubai + label: + en_US: Asia/Dubai + zh_Hans: 亚洲/迪拜 + pt_BR: Asia/Dubai + - value: Asia/Kolkata + label: + en_US: Asia/Kolkata + zh_Hans: 亚洲/加尔各答 + pt_BR: Asia/Kolkata + - value: Asia/Seoul + label: + en_US: Asia/Seoul + zh_Hans: 亚洲/首尔 + pt_BR: Asia/Seoul + - value: Asia/Singapore + label: + en_US: Asia/Singapore + zh_Hans: 亚洲/新加坡 + pt_BR: Asia/Singapore + - value: Europe/London + label: + en_US: Europe/London + zh_Hans: 欧洲/伦敦 + pt_BR: Europe/London + - value: Europe/Berlin + label: + en_US: Europe/Berlin + zh_Hans: 欧洲/柏林 + pt_BR: Europe/Berlin + - value: Europe/Moscow + label: + en_US: Europe/Moscow + zh_Hans: 欧洲/莫斯科 + pt_BR: Europe/Moscow + - value: Australia/Sydney + label: + en_US: Australia/Sydney + zh_Hans: 澳大利亚/悉尼 + pt_BR: Australia/Sydney + - value: Pacific/Auckland + label: + en_US: Pacific/Auckland + zh_Hans: 太平洋/奥克兰 + pt_BR: Pacific/Auckland + - value: Africa/Cairo + label: + en_US: Africa/Cairo + zh_Hans: 非洲/开罗 + pt_BR: Africa/Cairo diff --git a/api/core/tools/provider/builtin/time/tools/localtime_to_timestamp.py b/api/core/tools/provider/builtin/time/tools/localtime_to_timestamp.py new file mode 100644 index 0000000000000000000000000000000000000000..e16b732d0242db66c91adf32f41d60eca1ef6e15 --- /dev/null +++ b/api/core/tools/provider/builtin/time/tools/localtime_to_timestamp.py @@ -0,0 +1,44 @@ +from datetime import datetime +from typing import Any, Union + +import pytz + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.errors import ToolInvokeError +from core.tools.tool.builtin_tool import BuiltinTool + + +class LocaltimeToTimestampTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Convert localtime to timestamp + """ + localtime = tool_parameters.get("localtime") + timezone = tool_parameters.get("timezone", "Asia/Shanghai") + if not timezone: + timezone = None + time_format = "%Y-%m-%d %H:%M:%S" + + timestamp = self.localtime_to_timestamp(localtime, time_format, timezone) + if not timestamp: + return self.create_text_message(f"Invalid localtime: {localtime}") + + return self.create_text_message(f"{timestamp}") + + @staticmethod + def localtime_to_timestamp(localtime: str, time_format: str, local_tz=None) -> int | None: + try: + if local_tz is None: + local_tz = datetime.now().astimezone().tzinfo + if isinstance(local_tz, str): + local_tz = pytz.timezone(local_tz) + local_time = datetime.strptime(localtime, time_format) + localtime = local_tz.localize(local_time) + timestamp = int(localtime.timestamp()) + return timestamp + except Exception as e: + raise ToolInvokeError(str(e)) diff --git a/api/core/tools/provider/builtin/time/tools/localtime_to_timestamp.yaml b/api/core/tools/provider/builtin/time/tools/localtime_to_timestamp.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6a3b90595fd3fddb2986226bc4fc7db466b874b0 --- /dev/null +++ b/api/core/tools/provider/builtin/time/tools/localtime_to_timestamp.yaml @@ -0,0 +1,33 @@ +identity: + name: localtime_to_timestamp + author: zhuhao + label: + en_US: localtime to timestamp + zh_Hans: 获取时间戳 +description: + human: + en_US: A tool for localtime convert to timestamp + zh_Hans: 获取时间戳 + llm: A tool for localtime convert to timestamp +parameters: + - name: localtime + type: string + required: true + form: llm + label: + en_US: localtime + zh_Hans: 本地时间 + human_description: + en_US: localtime, such as 2024-1-1 0:0:0 + zh_Hans: 本地时间, 比如2024-1-1 0:0:0 + - name: timezone + type: string + required: false + form: llm + label: + en_US: Timezone + zh_Hans: 时区 + human_description: + en_US: Timezone, such as Asia/Shanghai + zh_Hans: 时区, 比如Asia/Shanghai + default: Asia/Shanghai diff --git a/api/core/tools/provider/builtin/time/tools/timestamp_to_localtime.py b/api/core/tools/provider/builtin/time/tools/timestamp_to_localtime.py new file mode 100644 index 0000000000000000000000000000000000000000..bcdd34fd4ec54d980de8057bdcb7d817ce5bb090 --- /dev/null +++ b/api/core/tools/provider/builtin/time/tools/timestamp_to_localtime.py @@ -0,0 +1,44 @@ +from datetime import datetime +from typing import Any, Union + +import pytz + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.errors import ToolInvokeError +from core.tools.tool.builtin_tool import BuiltinTool + + +class TimestampToLocaltimeTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Convert timestamp to localtime + """ + timestamp = tool_parameters.get("timestamp") + timezone = tool_parameters.get("timezone", "Asia/Shanghai") + if not timezone: + timezone = None + time_format = "%Y-%m-%d %H:%M:%S" + + locatime = self.timestamp_to_localtime(timestamp, timezone) + if not locatime: + return self.create_text_message(f"Invalid timestamp: {timestamp}") + + localtime_format = locatime.strftime(time_format) + + return self.create_text_message(f"{localtime_format}") + + @staticmethod + def timestamp_to_localtime(timestamp: int, local_tz=None) -> datetime | None: + try: + if local_tz is None: + local_tz = datetime.now().astimezone().tzinfo + if isinstance(local_tz, str): + local_tz = pytz.timezone(local_tz) + local_time = datetime.fromtimestamp(timestamp, local_tz) + return local_time + except Exception as e: + raise ToolInvokeError(str(e)) diff --git a/api/core/tools/provider/builtin/time/tools/timestamp_to_localtime.yaml b/api/core/tools/provider/builtin/time/tools/timestamp_to_localtime.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3794e717b4dc85301644bf25b4cc26ee9b58e05e --- /dev/null +++ b/api/core/tools/provider/builtin/time/tools/timestamp_to_localtime.yaml @@ -0,0 +1,33 @@ +identity: + name: timestamp_to_localtime + author: zhuhao + label: + en_US: Timestamp to localtime + zh_Hans: 时间戳转换 +description: + human: + en_US: A tool for timestamp convert to localtime + zh_Hans: 时间戳转换 + llm: A tool for timestamp convert to localtime +parameters: + - name: timestamp + type: number + required: true + form: llm + label: + en_US: Timestamp + zh_Hans: 时间戳 + human_description: + en_US: Timestamp + zh_Hans: 时间戳 + - name: timezone + type: string + required: false + form: llm + label: + en_US: Timezone + zh_Hans: 时区 + human_description: + en_US: Timezone, such as Asia/Shanghai + zh_Hans: 时区, 比如Asia/Shanghai + default: Asia/Shanghai diff --git a/api/core/tools/provider/builtin/time/tools/timezone_conversion.py b/api/core/tools/provider/builtin/time/tools/timezone_conversion.py new file mode 100644 index 0000000000000000000000000000000000000000..28e70db532852761be8c64debaf27f736cefe69b --- /dev/null +++ b/api/core/tools/provider/builtin/time/tools/timezone_conversion.py @@ -0,0 +1,48 @@ +from datetime import datetime +from typing import Any, Union + +import pytz + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.errors import ToolInvokeError +from core.tools.tool.builtin_tool import BuiltinTool + + +class TimezoneConversionTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Convert time to equivalent time zone + """ + current_time = tool_parameters.get("current_time") + current_timezone = tool_parameters.get("current_timezone", "Asia/Shanghai") + target_timezone = tool_parameters.get("target_timezone", "Asia/Tokyo") + target_time = self.timezone_convert(current_time, current_timezone, target_timezone) + if not target_time: + return self.create_text_message( + f"Invalid datatime and timezone: {current_time},{current_timezone},{target_timezone}" + ) + + return self.create_text_message(f"{target_time}") + + @staticmethod + def timezone_convert(current_time: str, source_timezone: str, target_timezone: str) -> str: + """ + Convert a time string from source timezone to target timezone. + """ + time_format = "%Y-%m-%d %H:%M:%S" + try: + # get source timezone + input_timezone = pytz.timezone(source_timezone) + # get target timezone + output_timezone = pytz.timezone(target_timezone) + local_time = datetime.strptime(current_time, time_format) + datetime_with_tz = input_timezone.localize(local_time) + # timezone convert + converted_datetime = datetime_with_tz.astimezone(output_timezone) + return converted_datetime.strftime(format=time_format) + except Exception as e: + raise ToolInvokeError(str(e)) diff --git a/api/core/tools/provider/builtin/time/tools/timezone_conversion.yaml b/api/core/tools/provider/builtin/time/tools/timezone_conversion.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4c221c2e512208fc52956190ae89b6fe89d0431c --- /dev/null +++ b/api/core/tools/provider/builtin/time/tools/timezone_conversion.yaml @@ -0,0 +1,44 @@ +identity: + name: timezone_conversion + author: zhuhao + label: + en_US: convert time to equivalent time zone + zh_Hans: 时区转换 +description: + human: + en_US: A tool to convert time to equivalent time zone + zh_Hans: 时区转换 + llm: A tool to convert time to equivalent time zone +parameters: + - name: current_time + type: string + required: true + form: llm + label: + en_US: current time + zh_Hans: 当前时间 + human_description: + en_US: current time, such as 2024-1-1 0:0:0 + zh_Hans: 当前时间, 比如2024-1-1 0:0:0 + - name: current_timezone + type: string + required: true + form: llm + label: + en_US: Current Timezone + zh_Hans: 当前时区 + human_description: + en_US: Current Timezone, such as Asia/Shanghai + zh_Hans: 当前时区, 比如Asia/Shanghai + default: Asia/Shanghai + - name: target_timezone + type: string + required: true + form: llm + label: + en_US: Target Timezone + zh_Hans: 目标时区 + human_description: + en_US: Target Timezone, such as Asia/Tokyo + zh_Hans: 目标时区, 比如Asia/Tokyo + default: Asia/Tokyo diff --git a/api/core/tools/provider/builtin/time/tools/weekday.py b/api/core/tools/provider/builtin/time/tools/weekday.py new file mode 100644 index 0000000000000000000000000000000000000000..b327e54e1710480b0685aa9bba24faf0bb59d5c5 --- /dev/null +++ b/api/core/tools/provider/builtin/time/tools/weekday.py @@ -0,0 +1,43 @@ +import calendar +from datetime import datetime +from typing import Any, Union + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class WeekdayTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Calculate the day of the week for a given date + """ + year = tool_parameters.get("year") + month = tool_parameters.get("month") + day = tool_parameters.get("day") + + date_obj = self.convert_datetime(year, month, day) + if not date_obj: + return self.create_text_message(f"Invalid date: Year {year}, Month {month}, Day {day}.") + + weekday_name = calendar.day_name[date_obj.weekday()] + month_name = calendar.month_name[month] + readable_date = f"{month_name} {date_obj.day}, {date_obj.year}" + return self.create_text_message(f"{readable_date} is {weekday_name}.") + + @staticmethod + def convert_datetime(year, month, day) -> datetime | None: + try: + # allowed range in datetime module + if not (year >= 1 and 1 <= month <= 12 and 1 <= day <= 31): + return None + + year = int(year) + month = int(month) + day = int(day) + return datetime(year, month, day) + except ValueError: + return None diff --git a/api/core/tools/provider/builtin/time/tools/weekday.yaml b/api/core/tools/provider/builtin/time/tools/weekday.yaml new file mode 100644 index 0000000000000000000000000000000000000000..481585e8c95c33c36dfca22a74d4d3c6718e5423 --- /dev/null +++ b/api/core/tools/provider/builtin/time/tools/weekday.yaml @@ -0,0 +1,42 @@ +identity: + name: weekday + author: Bowen Liang + label: + en_US: Weekday Calculator + zh_Hans: 星期几计算器 +description: + human: + en_US: A tool for calculating the weekday of a given date. + zh_Hans: 计算指定日期为星期几的工具。 + llm: A tool for calculating the weekday of a given date by year, month and day. +parameters: + - name: year + type: number + required: true + form: llm + label: + en_US: Year + zh_Hans: 年 + human_description: + en_US: Year + zh_Hans: 年 + - name: month + type: number + required: true + form: llm + label: + en_US: Month + zh_Hans: 月 + human_description: + en_US: Month + zh_Hans: 月 + - name: day + type: number + required: true + form: llm + label: + en_US: day + zh_Hans: 日 + human_description: + en_US: day + zh_Hans: 日 diff --git a/api/core/tools/provider/builtin/transcript/_assets/icon.svg b/api/core/tools/provider/builtin/transcript/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..83b0700fecbf30782d922a4e266946bbfd42dc83 --- /dev/null +++ b/api/core/tools/provider/builtin/transcript/_assets/icon.svg @@ -0,0 +1,11 @@ + + + + + + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/transcript/tools/transcript.py b/api/core/tools/provider/builtin/transcript/tools/transcript.py new file mode 100644 index 0000000000000000000000000000000000000000..ac7565d9eef5b8d4c052b745fed77b1d2e6d391c --- /dev/null +++ b/api/core/tools/provider/builtin/transcript/tools/transcript.py @@ -0,0 +1,81 @@ +from typing import Any, Union +from urllib.parse import parse_qs, urlparse + +from youtube_transcript_api import YouTubeTranscriptApi # type: ignore + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class YouTubeTranscriptTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Invoke the YouTube transcript tool + """ + try: + # Extract parameters with defaults + video_input = tool_parameters["video_id"] + language = tool_parameters.get("language") + output_format = tool_parameters.get("format", "text") + preserve_formatting = tool_parameters.get("preserve_formatting", False) + proxy = tool_parameters.get("proxy") + cookies = tool_parameters.get("cookies") + + # Extract video ID from URL if needed + video_id = self._extract_video_id(video_input) + + # Common kwargs for API calls + kwargs = {"proxies": {"https": proxy} if proxy else None, "cookies": cookies} + + try: + if language: + transcript_list = YouTubeTranscriptApi.list_transcripts(video_id, **kwargs) + try: + transcript = transcript_list.find_transcript([language]) + except: + # If requested language not found, try translating from English + transcript = transcript_list.find_transcript(["en"]).translate(language) + transcript_data = transcript.fetch() + else: + transcript_data = YouTubeTranscriptApi.get_transcript( + video_id, preserve_formatting=preserve_formatting, **kwargs + ) + + # Format output + formatter_class = { + "json": "JSONFormatter", + "pretty": "PrettyPrintFormatter", + "srt": "SRTFormatter", + "vtt": "WebVTTFormatter", + }.get(output_format) + + if formatter_class: + from youtube_transcript_api import formatters + + formatter = getattr(formatters, formatter_class)() + formatted_transcript = formatter.format_transcript(transcript_data) + else: + formatted_transcript = " ".join(entry["text"] for entry in transcript_data) + + return self.create_text_message(text=formatted_transcript) + + except Exception as e: + return self.create_text_message(text=f"Error getting transcript: {str(e)}") + + except Exception as e: + return self.create_text_message(text=f"Error processing request: {str(e)}") + + def _extract_video_id(self, video_input: str) -> str: + """ + Extract video ID from URL or return as-is if already an ID + """ + if "youtube.com" in video_input or "youtu.be" in video_input: + # Parse URL + parsed_url = urlparse(video_input) + if "youtube.com" in parsed_url.netloc: + return parse_qs(parsed_url.query)["v"][0] + else: # youtu.be + return parsed_url.path[1:] + return video_input # Assume it's already a video ID diff --git a/api/core/tools/provider/builtin/transcript/tools/transcript.yaml b/api/core/tools/provider/builtin/transcript/tools/transcript.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c654634a6c0be5f8b8907b1c754e3b9bf0f28c52 --- /dev/null +++ b/api/core/tools/provider/builtin/transcript/tools/transcript.yaml @@ -0,0 +1,101 @@ +identity: + name: free_youtube_transcript + author: Tao Wang + label: + en_US: Free YouTube Transcript API + zh_Hans: 免费获取 YouTube 转录 +description: + human: + en_US: Get transcript from a YouTube video for free. + zh_Hans: 免费获取 YouTube 视频的转录文案。 + llm: A tool for retrieving transcript from YouTube videos. +parameters: + - name: video_id + type: string + required: true + label: + en_US: Video ID/URL + zh_Hans: 视频ID + human_description: + en_US: Used to define the video from which the transcript will be fetched. You can find the id in the video url. For example - https://www.youtube.com/watch?v=video_id. + zh_Hans: 您要哪条视频的转录文案?您可以在视频链接中找到id。例如 - https://www.youtube.com/watch?v=video_id。 + llm_description: Used to define the video from which the transcript will be fetched. For example - https://www.youtube.com/watch?v=video_id. + form: llm + - name: language + type: string + required: false + label: + en_US: Language Code + zh_Hans: 语言 + human_description: + en_US: Language code (e.g. 'en', 'zh') for the transcript. + zh_Hans: 字幕语言代码(如'en'、'zh')。留空则自动选择。 + llm_description: Used to set the language for transcripts. + form: form + - name: format + type: select + required: false + default: text + options: + - value: text + label: + en_US: Plain Text + zh_Hans: 纯文本 + - value: json + label: + en_US: JSON Format + zh_Hans: JSON 格式 + - value: pretty + label: + en_US: Pretty Print Format + zh_Hans: 美化格式 + - value: srt + label: + en_US: SRT Format + zh_Hans: SRT 格式 + - value: vtt + label: + en_US: WebVTT Format + zh_Hans: WebVTT 格式 + label: + en_US: Output Format + zh_Hans: 输出格式 + human_description: + en_US: Format of the transcript output + zh_Hans: 字幕输出格式 + llm_description: The format to output the transcript in. Options are text (plain text), json (raw transcript data), srt (SubRip format), or vtt (WebVTT format) + form: form + - name: preserve_formatting + type: boolean + required: false + default: false + label: + en_US: Preserve Formatting + zh_Hans: 保留格式 + human_description: + en_US: Keep HTML formatting elements like (italics) and (bold) + zh_Hans: 保留HTML格式元素,如(斜体)和(粗体) + llm_description: Whether to preserve HTML formatting elements in the transcript text + form: form + - name: proxy + type: string + required: false + label: + en_US: HTTPS Proxy + zh_Hans: HTTPS 代理 + human_description: + en_US: HTTPS proxy URL (e.g. https://user:pass@domain:port) + zh_Hans: HTTPS 代理地址(如 https://user:pass@domain:port) + llm_description: HTTPS proxy to use for the request. Format should be https://user:pass@domain:port + form: form + - name: cookies + type: string + required: false + label: + en_US: Cookies File Path + zh_Hans: Cookies 文件路径 + human_description: + en_US: Path to cookies.txt file for accessing age-restricted videos + zh_Hans: 用于访问年龄限制视频的 cookies.txt 文件路径 + llm_description: Path to a cookies.txt file containing YouTube cookies, needed for accessing age-restricted videos + form: form diff --git a/api/core/tools/provider/builtin/transcript/transcript.py b/api/core/tools/provider/builtin/transcript/transcript.py new file mode 100644 index 0000000000000000000000000000000000000000..4fda1499882fe0cf1451424d60bc19b59f36b88f --- /dev/null +++ b/api/core/tools/provider/builtin/transcript/transcript.py @@ -0,0 +1,11 @@ +from typing import Any + +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class YouTubeTranscriptProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + """ + No credentials needed for YouTube Transcript API + """ + pass diff --git a/api/core/tools/provider/builtin/transcript/transcript.yaml b/api/core/tools/provider/builtin/transcript/transcript.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0786b454c35df33a19124d24e3d527b27319a989 --- /dev/null +++ b/api/core/tools/provider/builtin/transcript/transcript.yaml @@ -0,0 +1,13 @@ +identity: + author: Tao Wang + name: transcript + label: + en_US: Transcript + zh_Hans: Transcript + description: + en_US: Get transcripts from YouTube videos + zh_Hans: 获取 YouTube 视频的字幕/转录文本 + icon: icon.svg + tags: + - videos +credentials_for_provider: diff --git a/api/core/tools/provider/builtin/trello/_assets/icon.svg b/api/core/tools/provider/builtin/trello/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..f8e2bd47c0b818298a0dc6f426b11fa81bb6ed9b --- /dev/null +++ b/api/core/tools/provider/builtin/trello/_assets/icon.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/trello/tools/create_board.py b/api/core/tools/provider/builtin/trello/tools/create_board.py new file mode 100644 index 0000000000000000000000000000000000000000..5a61d2215789959c2da5fcb9455dea6110f43d05 --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/create_board.py @@ -0,0 +1,44 @@ +from typing import Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class CreateBoardTool(BuiltinTool): + """ + Tool for creating a new Trello board. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Union[str, int, bool]]) -> ToolInvokeMessage: + """ + Invoke the tool to create a new Trello board. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (dict[str, Union[str, int, bool]]): The parameters for the tool invocation. + + Returns: + ToolInvokeMessage: The result of the tool invocation. + """ + api_key = self.runtime.credentials.get("trello_api_key") + token = self.runtime.credentials.get("trello_api_token") + board_name = tool_parameters.get("name") + + if not (api_key and token and board_name): + return self.create_text_message("Missing required parameters: API key, token, or board name.") + + url = "https://api.trello.com/1/boards/" + query_params = {"name": board_name, "key": api_key, "token": token} + + try: + response = requests.post(url, params=query_params) + response.raise_for_status() + except requests.exceptions.RequestException as e: + return self.create_text_message("Failed to create board") + + board = response.json() + return self.create_text_message( + text=f"Board created successfully! Board name: {board['name']}, ID: {board['id']}" + ) diff --git a/api/core/tools/provider/builtin/trello/tools/create_board.yaml b/api/core/tools/provider/builtin/trello/tools/create_board.yaml new file mode 100644 index 0000000000000000000000000000000000000000..60dbab61f5ee5ce30a86dfc2e3cf0f574a994a75 --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/create_board.yaml @@ -0,0 +1,27 @@ +identity: + name: create_board + author: Yash Parmar + label: + en_US: Create Board + zh_Hans: 创建看板 + pt_BR: Criar Quadro +description: + human: + en_US: Creates a new Trello board with a specified name. This tool allows users to quickly add new boards to their Trello account, facilitating project organization and management. + zh_Hans: 使用指定的名称创建一个新的 Trello 看板。此工具允许用户快速向其 Trello 账户添加新的看板,促进项目组织和管理。 + pt_BR: Cria um novo quadro Trello com um nome especificado. Esta ferramenta permite que os usuários adicionem rapidamente novos quadros à sua conta Trello, facilitando a organização e gestão de projetos. + llm: Create a new Trello board using the specified name. This functionality simplifies the addition of boards, enhancing project organization and management within Trello. +parameters: + - name: name + type: string + required: true + label: + en_US: Board Name + zh_Hans: 看板名称 + pt_BR: Nome do Quadro + human_description: + en_US: The name for the new Trello board. This name helps in identifying and organizing your projects on Trello. + zh_Hans: 新 Trello 看板的名称。这个名称有助于在 Trello 上识别和组织您的项目。 + pt_BR: O nome para o novo quadro Trello. Este nome ajuda a identificar e organizar seus projetos no Trello. + llm_description: Specify the name for your new Trello board, aiding in project identification and organization within Trello. + form: llm diff --git a/api/core/tools/provider/builtin/trello/tools/create_list_on_board.py b/api/core/tools/provider/builtin/trello/tools/create_list_on_board.py new file mode 100644 index 0000000000000000000000000000000000000000..b32b0124dd31dae78eed9948ae5f0c6fd7403f5c --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/create_list_on_board.py @@ -0,0 +1,46 @@ +from typing import Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class CreateListOnBoardTool(BuiltinTool): + """ + Tool for creating a list on a Trello board by its ID. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Union[str, int, bool]]) -> ToolInvokeMessage: + """ + Invoke the tool to create a list on a Trello board by its ID. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (dict[str, Union[str, int, bool]]): The parameters for the tool invocation, + including the board ID and list name. + + Returns: + ToolInvokeMessage: The result of the tool invocation. + """ + api_key = self.runtime.credentials.get("trello_api_key") + token = self.runtime.credentials.get("trello_api_token") + board_id = tool_parameters.get("id") + list_name = tool_parameters.get("name") + + if not (api_key and token and board_id and list_name): + return self.create_text_message("Missing required parameters: API key, token, board ID, or list name.") + + url = f"https://api.trello.com/1/boards/{board_id}/lists" + params = {"name": list_name, "key": api_key, "token": token} + + try: + response = requests.post(url, params=params) + response.raise_for_status() + except requests.exceptions.RequestException as e: + return self.create_text_message("Failed to create list") + + new_list = response.json() + return self.create_text_message( + text=f"List '{new_list['name']}' created successfully with Id {new_list['id']} on board {board_id}." + ) diff --git a/api/core/tools/provider/builtin/trello/tools/create_list_on_board.yaml b/api/core/tools/provider/builtin/trello/tools/create_list_on_board.yaml new file mode 100644 index 0000000000000000000000000000000000000000..789b92437a3b3ec833491335516af0473d6284c1 --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/create_list_on_board.yaml @@ -0,0 +1,40 @@ +identity: + name: create_list_on_board + author: Yash Parmar + label: + en_US: Create List on Board + zh_Hans: 在看板上创建列表 + pt_BR: Criar Lista no Quadro +description: + human: + en_US: Creates a new list on a specified Trello board by providing the board's ID and the desired name for the list. Streamlines the process of organizing board content. + zh_Hans: 通过提供看板的 ID 和列表的所需名称,在指定的 Trello 看板上创建一个新列表。简化了组织看板内容的过程。 + pt_BR: Cria uma nova lista em um quadro Trello especificado, fornecendo o ID do quadro e o nome desejado para a lista. Facilita o processo de organização do conteúdo do quadro. + llm: Generate a new list within a Trello board by specifying the board's ID and a name for the list. Enhances board management by allowing quick additions of new lists. +parameters: + - name: id + type: string + required: true + label: + en_US: Board ID + zh_Hans: 看板 ID + pt_BR: ID do Quadro + human_description: + en_US: The unique identifier of the Trello board where the new list will be created. + zh_Hans: 新列表将被创建在其上的 Trello 看板的唯一标识符。 + pt_BR: O identificador único do quadro Trello onde a nova lista será criada. + llm_description: Input the ID of the Trello board to pinpoint where the new list should be added, ensuring correct placement. + form: llm + - name: name + type: string + required: true + label: + en_US: List Name + zh_Hans: 列表名称 + pt_BR: Nome da Lista + human_description: + en_US: The name for the new list to be created on the Trello board. + zh_Hans: 将在 Trello 看板上创建的新列表的名称。 + pt_BR: O nome para a nova lista que será criada no quadro Trello. + llm_description: Provide a name for the new list, defining its purpose or content focus, to facilitate board organization. + form: llm diff --git a/api/core/tools/provider/builtin/trello/tools/create_new_card_on_board.py b/api/core/tools/provider/builtin/trello/tools/create_new_card_on_board.py new file mode 100644 index 0000000000000000000000000000000000000000..e98efb81ca673e5e889da2920972281ec77813a8 --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/create_new_card_on_board.py @@ -0,0 +1,45 @@ +from typing import Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class CreateNewCardOnBoardTool(BuiltinTool): + """ + Tool for creating a new card on a Trello board. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Union[str, int, bool, None]]) -> ToolInvokeMessage: + """ + Invoke the tool to create a new card on a Trello board. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (dict[str, Union[str, int, bool, None]]): The parameters for the tool invocation, + including details for the new card. + + Returns: + ToolInvokeMessage: The result of the tool invocation. + """ + api_key = self.runtime.credentials.get("trello_api_key") + token = self.runtime.credentials.get("trello_api_token") + + # Ensure required parameters are present + if "name" not in tool_parameters or "idList" not in tool_parameters: + return self.create_text_message("Missing required parameters: name or idList.") + + url = "https://api.trello.com/1/cards" + params = {**tool_parameters, "key": api_key, "token": token} + + try: + response = requests.post(url, params=params) + response.raise_for_status() + new_card = response.json() + except requests.exceptions.RequestException as e: + return self.create_text_message("Failed to create card") + + return self.create_text_message( + text=f"New card '{new_card['name']}' created successfully with ID {new_card['id']}." + ) diff --git a/api/core/tools/provider/builtin/trello/tools/create_new_card_on_board.yaml b/api/core/tools/provider/builtin/trello/tools/create_new_card_on_board.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9953af718ddd581ffe42fc0a84b034b4165ea667 --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/create_new_card_on_board.yaml @@ -0,0 +1,145 @@ +identity: + name: create_new_card_on_board + author: Yash Parmar + label: + en_US: Create New Card on Board + zh_Hans: 在看板上创建新卡片 + pt_BR: Criar Novo Cartão no Quadro +description: + human: + en_US: Creates a new card on a Trello board with specified details like name, description, list ID, and other optional parameters. Facilitates task addition and project management within Trello. + zh_Hans: 用指定的详情(如名称、描述、列表 ID 和其他可选参数)在 Trello 看板上创建一个新卡片。便于在 Trello 中添加任务和管理项目。 + pt_BR: Cria um novo cartão em um quadro Trello com detalhes especificados, como nome, descrição, ID da lista e outros parâmetros opcionais. Facilita a adição de tarefas e a gestão de projetos dentro do Trello. + llm: Initiate a new card on a Trello board by specifying essential details such as the card's name, description, and the list it belongs to, among other settings. Streamlines project task additions and organizational workflows. +parameters: + - name: name + type: string + required: true + label: + en_US: Card Name + zh_Hans: 卡片名称 + pt_BR: Nome do Cartão + human_description: + en_US: The name for the new card. Acts as the primary identifier and summary of the card's purpose. + zh_Hans: 新卡片的名称。作为卡片目的的主要标识和总结。 + pt_BR: O nome para o novo cartão. Funciona como o identificador principal e resumo do propósito do cartão. + llm_description: Provide a concise, descriptive name for the card, outlining its main focus or task. + form: llm + # Include additional parameters like desc, pos, due, idList, etc., following the same pattern. + - name: desc + type: string + required: false + label: + en_US: Card Description + zh_Hans: 卡片描述 + pt_BR: Descrição do Cartão + human_description: + en_US: Optional. A brief description of the card's purpose or contents. + zh_Hans: 可选。卡片目的或内容的简要描述。 + pt_BR: Opcional. Uma breve descrição do propósito ou conteúdo do cartão. + llm_description: Add a brief description to the card to provide context or additional information about its purpose. + form: llm + - name: pos + type: string + required: false + label: + en_US: Position + zh_Hans: 位置 + pt_BR: Posição + human_description: + en_US: Optional. The position of the card in the list. Can be 'top', 'bottom', or a positive number. + zh_Hans: 可选。卡片在列表中的位置。可以是“top”、“bottom” 或正数。 + pt_BR: Opcional. A posição do cartão na lista. Pode ser 'top', 'bottom' ou um número positivo. + llm_description: Specify the position of the card within the list, either at the top, bottom, or a specific numerical index. + form: llm + - name: due + type: string + required: false + label: + en_US: Due Date + zh_Hans: 截止日期 + pt_BR: Data de Vencimento + human_description: + en_US: Optional. The due date for the card in the format 'MM/DD/YYYY'. + zh_Hans: 可选。卡片的截止日期,格式为“MM/DD/YYYY”。 + pt_BR: Opcional. A data de vencimento do cartão no formato 'MM/DD/YYYY'. + llm_description: Set a due date for the card to establish a deadline for completion or action. + form: llm + - name: start + type: string + required: false + label: + en_US: Start Date + zh_Hans: 开始日期 + pt_BR: Data de Início + human_description: + en_US: Optional. The start date for the card in the format 'MM/DD/YYYY'. + zh_Hans: 可选。卡片的开始日期,格式为“MM/DD/YYYY”。 + pt_BR: Opcional. A data de início do cartão no formato 'MM/DD/YYYY'. + llm_description: Specify a start date for the card to mark the beginning of a task or project phase. + form: llm + - name: dueComplete + type: boolean + required: false + label: + en_US: Due Complete + zh_Hans: 截止日期已完成 + pt_BR: Vencimento Concluído + human_description: + en_US: Optional. Set to true if the due date has been completed, or false if it is pending. + zh_Hans: 可选。如果截止日期已完成,则设置为 true;如果尚未完成,则设置为 false。 + pt_BR: Opcional. Defina como true se a data de vencimento foi concluída, ou como false se estiver pendente. + llm_description: Indicate whether the due date for the card has been marked as complete or is still pending. + form: llm + - name: idList + type: string + required: true + label: + en_US: List ID + zh_Hans: 列表 ID + pt_BR: ID da Lista + human_description: + en_US: The unique identifier of the list where the card will be added. + zh_Hans: 卡片将被添加到的列表的唯一标识符。 + pt_BR: O identificador único da lista onde o cartão será adicionado. + llm_description: Input the ID of the list where the card should be placed, ensuring it is added to the correct list. + form: llm + - name: idMembers + type: string + required: false + label: + en_US: Member IDs + zh_Hans: 成员 ID + pt_BR: IDs de Membros + human_description: + en_US: Optional. The IDs of members to assign to the card. + zh_Hans: 可选。要分配给卡片的成员的 ID。 + pt_BR: Opcional. Os IDs dos membros a serem atribuídos ao cartão. + llm_description: Specify the IDs of members to assign to the card, allowing for task delegation or collaboration. + form: llm + - name: idLabels + type: string + required: false + label: + en_US: Label IDs + zh_Hans: 标签 ID + pt_BR: IDs de Etiquetas + human_description: + en_US: Optional. The IDs of labels to assign to the card. + zh_Hans: 可选。要分配给卡片的标签的 ID。 + pt_BR: Opcional. Os IDs das etiquetas a serem atribuídos ao cartão. + llm_description: Assign specific labels to the card by providing their IDs, aiding in visual categorization or prioritization. + form: llm + - name: urlSource + type: string + required: false + label: + en_US: Source URL + zh_Hans: 来源 URL + pt_BR: URL de Origem + human_description: + en_US: Optional. The URL to attach as the card's source. + zh_Hans: 可选。要附加为卡片来源的 URL。 + pt_BR: Opcional. O URL a ser anexado como a fonte do cartão. + llm_description: Provide a URL to serve as the source reference for the card, linking to external resources or documents. + form: llm diff --git a/api/core/tools/provider/builtin/trello/tools/delete_board.py b/api/core/tools/provider/builtin/trello/tools/delete_board.py new file mode 100644 index 0000000000000000000000000000000000000000..7fc9d1f13c2664015cae9f3337adce4547559f71 --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/delete_board.py @@ -0,0 +1,41 @@ +from typing import Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class DeleteBoardTool(BuiltinTool): + """ + Tool for deleting a Trello board by ID. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Union[str, int, bool]]) -> ToolInvokeMessage: + """ + Invoke the tool to delete a Trello board by its ID. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (dict[str, Union[str, int, bool]]): The parameters for the tool invocation, + including the board ID. + + Returns: + ToolInvokeMessage: The result of the tool invocation. + """ + api_key = self.runtime.credentials.get("trello_api_key") + token = self.runtime.credentials.get("trello_api_token") + board_id = tool_parameters.get("boardId") + + if not (api_key and token and board_id): + return self.create_text_message("Missing required parameters: API key, token, or board ID.") + + url = f"https://api.trello.com/1/boards/{board_id}?key={api_key}&token={token}" + + try: + response = requests.delete(url) + response.raise_for_status() + except requests.exceptions.RequestException as e: + return self.create_text_message("Failed to delete board") + + return self.create_text_message(text=f"Board with ID {board_id} deleted successfully.") diff --git a/api/core/tools/provider/builtin/trello/tools/delete_board.yaml b/api/core/tools/provider/builtin/trello/tools/delete_board.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f043e78870d062cfb3a05ea64f5933225cf94a7d --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/delete_board.yaml @@ -0,0 +1,27 @@ +identity: + name: delete_board + author: Yash Parmar + label: + en_US: Delete Board + zh_Hans: 删除看板 + pt_BR: Excluir Quadro +description: + human: + en_US: Deletes a Trello board using its unique ID. This tool allows for the removal of boards that are no longer needed, ensuring a tidy workspace. + zh_Hans: 使用其唯一 ID 删除 Trello 看板。此工具允许删除不再需要的看板,确保工作区整洁。 + pt_BR: Exclui um quadro Trello usando seu ID único. Esta ferramenta permite a remoção de quadros que não são mais necessários, garantindo um espaço de trabalho organizado. + llm: Remove a Trello board by specifying its ID. This functionality is helpful for cleaning up unnecessary boards from your Trello account. +parameters: + - name: boardId + type: string + required: true + label: + en_US: Board ID + zh_Hans: 看板 ID + pt_BR: ID do Quadro + human_description: + en_US: The unique identifier for the Trello board you wish to delete. This ensures the specific board is accurately targeted for deletion. + zh_Hans: 您希望删除的 Trello 看板的唯一标识符。这确保了准确地针对特定看板进行删除。 + pt_BR: O identificador único para o quadro Trello que você deseja excluir. Isso garante que o quadro específico seja precisamente direcionado para exclusão. + llm_description: Enter the ID of the Trello board you want to remove. This ID is essential to identify the board precisely and perform the deletion. + form: llm diff --git a/api/core/tools/provider/builtin/trello/tools/delete_card.py b/api/core/tools/provider/builtin/trello/tools/delete_card.py new file mode 100644 index 0000000000000000000000000000000000000000..1de98d639ebb7d996cb8065e9d2f312e92260380 --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/delete_card.py @@ -0,0 +1,41 @@ +from typing import Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class DeleteCardByIdTool(BuiltinTool): + """ + Tool for deleting a Trello card by its ID. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Union[str, int, bool]]) -> ToolInvokeMessage: + """ + Invoke the tool to delete a Trello card by its ID. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (dict[str, Union[str, int, bool]]): The parameters for the tool invocation, + including the card ID. + + Returns: + ToolInvokeMessage: The result of the tool invocation. + """ + api_key = self.runtime.credentials.get("trello_api_key") + token = self.runtime.credentials.get("trello_api_token") + card_id = tool_parameters.get("id") + + if not (api_key and token and card_id): + return self.create_text_message("Missing required parameters: API key, token, or card ID.") + + url = f"https://api.trello.com/1/cards/{card_id}?key={api_key}&token={token}" + + try: + response = requests.delete(url) + response.raise_for_status() + except requests.exceptions.RequestException as e: + return self.create_text_message("Failed to delete card") + + return self.create_text_message(text=f"Card with ID {card_id} has been successfully deleted.") diff --git a/api/core/tools/provider/builtin/trello/tools/delete_card.yaml b/api/core/tools/provider/builtin/trello/tools/delete_card.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8898ef1bde3680bd561e24a427f78576ed7121dc --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/delete_card.yaml @@ -0,0 +1,27 @@ +identity: + name: delete_card_by_id + author: Yash Parmar + label: + en_US: Delete Card by ID + zh_Hans: 通过 ID 删除卡片 + pt_BR: Deletar Cartão por ID +description: + human: + en_US: Deletes a Trello card using its unique ID. This tool facilitates the removal of cards that are no longer needed, maintaining an organized board. + zh_Hans: 使用其唯一 ID 删除 Trello 卡片。此工具便于删除不再需要的卡片,保持看板的有序。 + pt_BR: Exclui um cartão Trello usando seu ID único. Esta ferramenta facilita a remoção de cartões que não são mais necessários, mantendo um quadro organizado. + llm: Remove a specific Trello card by providing its ID. Ideal for cleaning up and organizing your Trello boards by eliminating unwanted cards. +parameters: + - name: id + type: string + required: true + label: + en_US: Card ID + zh_Hans: 卡片 ID + pt_BR: ID do Cartão + human_description: + en_US: The unique identifier of the Trello card you wish to delete. This ensures the precise card is removed. + zh_Hans: 您希望删除的 Trello 卡片的唯一标识符。这确保了精确移除特定卡片。 + pt_BR: O identificador único do cartão Trello que você deseja excluir. Isso garante que o cartão exato seja removido. + llm_description: Input the ID of the Trello card targeted for deletion to ensure accurate and specific removal. + form: llm diff --git a/api/core/tools/provider/builtin/trello/tools/fetch_all_boards.py b/api/core/tools/provider/builtin/trello/tools/fetch_all_boards.py new file mode 100644 index 0000000000000000000000000000000000000000..0c5ed9ea8533ff0a831ffa68602c1b03560d45d1 --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/fetch_all_boards.py @@ -0,0 +1,50 @@ +from typing import Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class FetchAllBoardsTool(BuiltinTool): + """ + Tool for fetching all boards from Trello. + """ + + def _invoke( + self, user_id: str, tool_parameters: dict[str, Union[str, int, bool]] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Invoke the fetch all boards tool. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (dict[str, Union[str, int, bool]]): The parameters for the tool invocation. + + Returns: + Union[ToolInvokeMessage, List[ToolInvokeMessage]]: The result of the tool invocation. + """ + api_key = self.runtime.credentials.get("trello_api_key") + token = self.runtime.credentials.get("trello_api_token") + + if not (api_key and token): + return self.create_text_message("Missing Trello API key or token in credentials.") + + # Including board filter in the request if provided + board_filter = tool_parameters.get("boards", "open") + url = f"https://api.trello.com/1/members/me/boards?filter={board_filter}&key={api_key}&token={token}" + + try: + response = requests.get(url) + response.raise_for_status() # Raises stored HTTPError, if one occurred. + except requests.exceptions.RequestException as e: + return self.create_text_message("Failed to fetch boards") + + boards = response.json() + + if not boards: + return self.create_text_message("No boards found in Trello.") + + # Creating a string with both board names and IDs + boards_info = ", ".join([f"{board['name']} (ID: {board['id']})" for board in boards]) + return self.create_text_message(text=f"Boards: {boards_info}") diff --git a/api/core/tools/provider/builtin/trello/tools/fetch_all_boards.yaml b/api/core/tools/provider/builtin/trello/tools/fetch_all_boards.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d0ac4beaaa723a782cdcc3b21cdc4fd43a5d553e --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/fetch_all_boards.yaml @@ -0,0 +1,28 @@ +identity: + name: fetch_all_boards + author: Yash Parmar + label: + en_US: Fetch All Boards + zh_Hans: 获取所有看板 + pt_BR: Buscar Todos os Quadros +description: + human: + en_US: Retrieves all the Trello boards associated with the user's account. This tool provides a quick overview of all open boards, aiding in efficient project management and organization. + zh_Hans: 检索与用户账户关联的所有 Trello 看板。该工具提供了所有打开的看板的快速概览,有助于高效的项目管理和组织。 + pt_BR: Recupera todos os quadros do Trello associados à conta do usuário. Esta ferramenta oferece uma visão geral rápida de todos os quadros abertos, auxiliando na gestão e organização eficiente do projeto. + llm: This tool fetches all Trello boards linked to the user's account, offering a swift snapshot of open boards to streamline project management and organization tasks. +parameters: + - name: boards + type: string + required: false + default: open + label: + en_US: Boards filter + zh_Hans: 看板过滤器 + pt_BR: Filtro de quadros + human_description: + en_US: Specifies the type of boards to retrieve. Default is 'open', fetching all open boards. Other options include 'closed', 'members', 'organization', etc. + zh_Hans: 指定要检索的看板类型。默认为“open”,获取所有打开的看板。其他选项包括“closed”,“members”,“organization”等。 + pt_BR: Especifica o tipo de quadros a serem recuperados. O padrão é 'open', buscando todos os quadros abertos. Outras opções incluem 'closed', 'members', 'organization', etc. + llm_description: Determines the category of boards to be displayed, with 'open' as the default setting to show all open boards. Variants like 'closed', 'members', and 'organization' are also selectable. + form: llm diff --git a/api/core/tools/provider/builtin/trello/tools/get_board_actions.py b/api/core/tools/provider/builtin/trello/tools/get_board_actions.py new file mode 100644 index 0000000000000000000000000000000000000000..cabc7ce09359d54d1166255ebdfdfaa22c25f421 --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/get_board_actions.py @@ -0,0 +1,45 @@ +from typing import Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class GetBoardActionsTool(BuiltinTool): + """ + Tool for retrieving actions for a Trello board by its ID. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Union[str, int, bool]]) -> ToolInvokeMessage: + """ + Invoke the tool to retrieve actions for a Trello board by its ID. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (dict[str, Union[str, int, bool]]): The parameters for the tool invocation, + including the board ID. + + Returns: + ToolInvokeMessage: The result of the tool invocation. + """ + api_key = self.runtime.credentials.get("trello_api_key") + token = self.runtime.credentials.get("trello_api_token") + board_id = tool_parameters.get("boardId") + + if not (api_key and token and board_id): + return self.create_text_message("Missing required parameters: API key, token, or board ID.") + + url = f"https://api.trello.com/1/boards/{board_id}/actions?key={api_key}&token={token}" + + try: + response = requests.get(url) + response.raise_for_status() + actions = response.json() + except requests.exceptions.RequestException as e: + return self.create_text_message("Failed to retrieve board actions") + + actions_summary = "\n".join( + [f"{action['type']}: {action.get('data', {}).get('text', 'No details available')}" for action in actions] + ) + return self.create_text_message(text=f"Actions for Board ID {board_id}:\n{actions_summary}") diff --git a/api/core/tools/provider/builtin/trello/tools/get_board_actions.yaml b/api/core/tools/provider/builtin/trello/tools/get_board_actions.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1ba89f9e44abbd35af565f479ac253d6e04da47e --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/get_board_actions.yaml @@ -0,0 +1,27 @@ +identity: + name: get_board_actions + author: Yash Parmar + label: + en_US: Get Board Actions + zh_Hans: 获取看板操作 + pt_BR: Obter Ações do Quadro +description: + human: + en_US: Retrieves a list of actions (such as updates, movements, and comments) for a Trello board by its ID. This tool provides insights into the board's activity history. + zh_Hans: 通过其 ID 为 Trello 看板检索操作列表(如更新、移动和评论)。此工具提供了看板活动历史的见解。 + pt_BR: Recupera uma lista de ações (como atualizações, movimentos e comentários) para um quadro Trello pelo seu ID. Esta ferramenta oferece insights sobre o histórico de atividades do quadro. + llm: Fetch the sequence of actions performed on a Trello board, such as card updates, movements, and comments, by providing the board's ID. Offers a historical view of board activities. +parameters: + - name: boardId + type: string + required: true + label: + en_US: Board ID + zh_Hans: 看板 ID + pt_BR: ID do Quadro + human_description: + en_US: The unique identifier of the Trello board for which you want to retrieve actions. It targets the specific board to fetch its activity log. + zh_Hans: 您想要检索操作的 Trello 看板的唯一标识符。它定位特定的看板以获取其活动日志。 + pt_BR: O identificador único do quadro Trello para o qual você deseja recuperar ações. Direciona especificamente para o quadro para buscar seu registro de atividades. + llm_description: Input the ID of the Trello board to access its detailed action history, including all updates, comments, and movements related to the board. + form: llm diff --git a/api/core/tools/provider/builtin/trello/tools/get_board_by_id.py b/api/core/tools/provider/builtin/trello/tools/get_board_by_id.py new file mode 100644 index 0000000000000000000000000000000000000000..fe42cd9c5cbf863f96b4b1715f18c32880e851f6 --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/get_board_by_id.py @@ -0,0 +1,66 @@ +from typing import Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class GetBoardByIdTool(BuiltinTool): + """ + Tool for retrieving detailed information about a Trello board by its ID. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Union[str, int, bool]]) -> ToolInvokeMessage: + """ + Invoke the tool to retrieve a Trello board by its ID. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (dict[str, Union[str, int, bool]]): The parameters for the tool invocation, + including the board ID. + + Returns: + ToolInvokeMessage: The result of the tool invocation. + """ + api_key = self.runtime.credentials.get("trello_api_key") + token = self.runtime.credentials.get("trello_api_token") + board_id = tool_parameters.get("boardId") + + if not (api_key and token and board_id): + return self.create_text_message("Missing required parameters: API key, token, or board ID.") + + url = f"https://api.trello.com/1/boards/{board_id}?key={api_key}&token={token}" + + try: + response = requests.get(url) + response.raise_for_status() + board = response.json() + board_details = self.format_board_details(board) + except requests.exceptions.RequestException as e: + return self.create_text_message("Failed to retrieve board") + + return self.create_text_message(text=board_details) + + def format_board_details(self, board: dict) -> str: + """ + Format the board details into a human-readable string. + + Args: + board (dict): The board information as a dictionary. + + Returns: + str: Formatted board details. + """ + details = ( + f"Board Name: {board['name']}\n" + f"Board ID: {board['id']}\n" + f"Description: {board['desc'] or 'No description provided.'}\n" + f"Status: {'Closed' if board['closed'] else 'Open'}\n" + f"Organization ID: {board['idOrganization'] or 'Not part of an organization.'}\n" + f"URL: {board['url']}\n" + f"Short URL: {board['shortUrl']}\n" + f"Permission Level: {board['prefs']['permissionLevel']}\n" + f"Background Color: {board['prefs']['backgroundColor']}" + ) + return details diff --git a/api/core/tools/provider/builtin/trello/tools/get_board_by_id.yaml b/api/core/tools/provider/builtin/trello/tools/get_board_by_id.yaml new file mode 100644 index 0000000000000000000000000000000000000000..45c93006ba441433d0c7f2b7e8875b0bd1197740 --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/get_board_by_id.yaml @@ -0,0 +1,27 @@ +identity: + name: get_board_by_id + author: Yash Parmar + label: + en_US: Get Board by ID + zh_Hans: 通过 ID 获取看板 + pt_BR: Obter Quadro por ID +description: + human: + en_US: Retrieves detailed information about a specific Trello board using its unique ID. This tool enables users to quickly access board details without navigating through the Trello interface. + zh_Hans: 使用其唯一 ID 检索有关特定 Trello 看板的详细信息。此工具使用户能够快速访问看板详情,无需通过 Trello 界面导航。 + pt_BR: Recupera informações detalhadas sobre um quadro Trello específico usando seu ID único. Esta ferramenta permite que os usuários acessem rapidamente os detalhes do quadro sem navegar pela interface do Trello. + llm: Access details of a Trello board by providing its ID. This tool offers a direct way to view board information, simplifying the process of managing and reviewing Trello boards. +parameters: + - name: boardId + type: string + required: true + label: + en_US: Board ID + zh_Hans: 看板 ID + pt_BR: ID do Quadro + human_description: + en_US: The unique identifier for the Trello board you wish to retrieve. This ID enables precise targeting and fetching of the board's details. + zh_Hans: 您希望检索的 Trello 看板的唯一标识符。此 ID 使能够准确定位和获取看板的详细信息。 + pt_BR: O identificador único do quadro Trello que você deseja recuperar. Este ID permite o direcionamento preciso e a obtenção dos detalhes do quadro. + llm_description: Input the ID of the Trello board to get its details. This unique ID ensures accurate retrieval of information about the specified board. + form: llm diff --git a/api/core/tools/provider/builtin/trello/tools/get_board_cards.py b/api/core/tools/provider/builtin/trello/tools/get_board_cards.py new file mode 100644 index 0000000000000000000000000000000000000000..ff2b1221e767de996bdf5b4c745cb935b79be7db --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/get_board_cards.py @@ -0,0 +1,43 @@ +from typing import Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class GetBoardCardsTool(BuiltinTool): + """ + Tool for retrieving cards on a Trello board by its ID. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Union[str, int, bool]]) -> ToolInvokeMessage: + """ + Invoke the tool to retrieve cards on a Trello board by its ID. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (dict[str, Union[str, int, bool]]): The parameters for the tool invocation, + including the board ID. + + Returns: + ToolInvokeMessage: The result of the tool invocation. + """ + api_key = self.runtime.credentials.get("trello_api_key") + token = self.runtime.credentials.get("trello_api_token") + board_id = tool_parameters.get("boardId") + + if not (api_key and token and board_id): + return self.create_text_message("Missing required parameters: API key, token, or board ID.") + + url = f"https://api.trello.com/1/boards/{board_id}/cards?key={api_key}&token={token}" + + try: + response = requests.get(url) + response.raise_for_status() + cards = response.json() + except requests.exceptions.RequestException as e: + return self.create_text_message("Failed to retrieve board cards") + + cards_summary = "\n".join([f"{card['name']} (ID: {card['id']})" for card in cards]) + return self.create_text_message(text=f"Cards for Board ID {board_id}:\n{cards_summary}") diff --git a/api/core/tools/provider/builtin/trello/tools/get_board_cards.yaml b/api/core/tools/provider/builtin/trello/tools/get_board_cards.yaml new file mode 100644 index 0000000000000000000000000000000000000000..852ea278af341cf9d6d9070f971315ac56ea500a --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/get_board_cards.yaml @@ -0,0 +1,27 @@ +identity: + name: get_board_cards + author: Yash Parmar + label: + en_US: Get Board Cards + zh_Hans: 获取看板卡片 + pt_BR: Obter Cartões do Quadro +description: + human: + en_US: Retrieves all cards present on a specific Trello board by its ID, providing a list of card names and their IDs. Useful for managing and organizing project tasks. + zh_Hans: 通过其 ID 检索特定 Trello 看板上的所有卡片,提供卡片名称及其 ID 的列表。用于管理和组织项目任务。 + pt_BR: Recupera todos os cartões presentes em um quadro Trello específico pelo seu ID, fornecendo uma lista dos nomes dos cartões e seus IDs. Útil para gerenciar e organizar tarefas de projetos. + llm: Obtain a list of all cards on a specific Trello board by entering the board's ID. This tool helps in quickly assessing the tasks or items associated with the board. +parameters: + - name: boardId + type: string + required: true + label: + en_US: Board ID + zh_Hans: 看板 ID + pt_BR: ID do Quadro + human_description: + en_US: The unique identifier of the Trello board from which you want to retrieve cards. It specifies the exact board to gather card details from. + zh_Hans: 您想要从中检索卡片的 Trello 看板的唯一标识符。它指定了要从中收集卡片详细信息的确切看板。 + pt_BR: O identificador único do quadro Trello do qual você deseja recuperar os cartões. Especifica o quadro exato para obter detalhes dos cartões. + llm_description: Input the ID of the Trello board to fetch its cards, allowing for a detailed overview of the board's contents. + form: llm diff --git a/api/core/tools/provider/builtin/trello/tools/get_filterd_board_cards.py b/api/core/tools/provider/builtin/trello/tools/get_filterd_board_cards.py new file mode 100644 index 0000000000000000000000000000000000000000..3d7f9f4ad1c99641a30c6af42766eb2319d4a217 --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/get_filterd_board_cards.py @@ -0,0 +1,46 @@ +from typing import Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class GetFilteredBoardCardsTool(BuiltinTool): + """ + Tool for retrieving filtered cards on a Trello board by its ID and a specified filter. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Union[str, int, bool]]) -> ToolInvokeMessage: + """ + Invoke the tool to retrieve filtered cards on a Trello board by its ID and filter. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (dict[str, Union[str, int, bool]]): The parameters for the tool invocation, + including the board ID and filter. + + Returns: + ToolInvokeMessage: The result of the tool invocation. + """ + api_key = self.runtime.credentials.get("trello_api_key") + token = self.runtime.credentials.get("trello_api_token") + board_id = tool_parameters.get("boardId") + filter = tool_parameters.get("filter") + + if not (api_key and token and board_id and filter): + return self.create_text_message("Missing required parameters: API key, token, board ID, or filter.") + + url = f"https://api.trello.com/1/boards/{board_id}/cards/{filter}?key={api_key}&token={token}" + + try: + response = requests.get(url) + response.raise_for_status() + filtered_cards = response.json() + except requests.exceptions.RequestException as e: + return self.create_text_message("Failed to retrieve filtered cards") + + card_details = "\n".join([f"{card['name']} (ID: {card['id']})" for card in filtered_cards]) + return self.create_text_message( + text=f"Filtered Cards for Board ID {board_id} with Filter '{filter}':\n{card_details}" + ) diff --git a/api/core/tools/provider/builtin/trello/tools/get_filterd_board_cards.yaml b/api/core/tools/provider/builtin/trello/tools/get_filterd_board_cards.yaml new file mode 100644 index 0000000000000000000000000000000000000000..390595645771e4c7ab851a9a11ad03284297d556 --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/get_filterd_board_cards.yaml @@ -0,0 +1,40 @@ +identity: + name: get_filtered_board_cards + author: Yash Parmar + label: + en_US: Get Filtered Board Cards + zh_Hans: 获取筛选的看板卡片 + pt_BR: Obter Cartões Filtrados do Quadro +description: + human: + en_US: Retrieves cards from a Trello board using a specified filter and the board's ID. Filters include options like 'all', 'open', 'closed', 'none', and 'visible', allowing for tailored views of board content. + zh_Hans: 使用指定的过滤器和看板的 ID 从 Trello 看板检索卡片。过滤器包括 'all', 'open', 'closed', 'none' 和 'visible' 等选项,允许对看板内容进行定制查看。 + pt_BR: Recupera cartões de um quadro Trello usando um filtro especificado e o ID do quadro. Os filtros incluem opções como 'all', 'open', 'closed', 'none' e 'visible', permitindo visualizações personalizadas do conteúdo do quadro. + llm: Access cards on a Trello board through specific filters such as 'all', 'open', 'closed', 'none', and 'visible' by providing the board's ID. This feature enables focused examination of the board's cards. +parameters: + - name: boardId + type: string + required: true + label: + en_US: Board ID + zh_Hans: 看板 ID + pt_BR: ID do Quadro + human_description: + en_US: The unique identifier for the Trello board from which to retrieve the filtered cards. + zh_Hans: 用于检索筛选卡片的 Trello 看板的唯一标识符。 + pt_BR: O identificador único do quadro Trello do qual os cartões filtrados serão recuperados. + llm_description: Enter the Trello board's ID to specify from which board to fetch the cards using the filter. + form: llm + - name: filter + type: string + required: true + label: + en_US: Filter + zh_Hans: 过滤器 + pt_BR: Filtro + human_description: + en_US: The filter to apply when retrieving cards. Valid values are 'all', 'open', 'closed', 'none', and 'visible'. + zh_Hans: 检索卡片时应用的过滤器。有效值为 'all', 'open', 'closed', 'none', 和 'visible'。 + pt_BR: O filtro a ser aplicado ao recuperar cartões. Os valores válidos são 'all', 'open', 'closed', 'none' e 'visible'. + llm_description: Specify the filter for card retrieval. Choose from 'all', 'open', 'closed', 'none', or 'visible' to control which cards are fetched. + form: llm diff --git a/api/core/tools/provider/builtin/trello/tools/get_lists_on_board.py b/api/core/tools/provider/builtin/trello/tools/get_lists_on_board.py new file mode 100644 index 0000000000000000000000000000000000000000..ccf404068f225e0b7c00c1a80f74b133f2ce2fe4 --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/get_lists_on_board.py @@ -0,0 +1,43 @@ +from typing import Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class GetListsFromBoardTool(BuiltinTool): + """ + Tool for retrieving all lists from a specified Trello board by its ID. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Union[str, int, bool]]) -> ToolInvokeMessage: + """ + Invoke the tool to get all lists from a specified Trello board. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (dict[str, Union[str, int, bool]]): The parameters for the tool invocation, + including the board ID. + + Returns: + ToolInvokeMessage: The result of the tool invocation. + """ + api_key = self.runtime.credentials.get("trello_api_key") + token = self.runtime.credentials.get("trello_api_token") + board_id = tool_parameters.get("boardId") + + if not (api_key and token and board_id): + return self.create_text_message("Missing required parameters: API key, token, or board ID.") + + url = f"https://api.trello.com/1/boards/{board_id}/lists?key={api_key}&token={token}" + + try: + response = requests.get(url) + response.raise_for_status() + lists = response.json() + except requests.exceptions.RequestException as e: + return self.create_text_message("Failed to retrieve lists") + + lists_info = "\n".join([f"{list['name']} (ID: {list['id']})" for list in lists]) + return self.create_text_message(text=f"Lists on Board ID {board_id}:\n{lists_info}") diff --git a/api/core/tools/provider/builtin/trello/tools/get_lists_on_board.yaml b/api/core/tools/provider/builtin/trello/tools/get_lists_on_board.yaml new file mode 100644 index 0000000000000000000000000000000000000000..31028a80404de35a526a3cd3447ade7837679f10 --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/get_lists_on_board.yaml @@ -0,0 +1,27 @@ +identity: + name: get_lists_from_board + author: Yash Parmar + label: + en_US: Get Lists from Board + zh_Hans: 获取看板的列表 + pt_BR: Obter Listas do Quadro +description: + human: + en_US: Retrieves all lists from a specified Trello board by its ID, providing an overview of the board's organization and current phases or categories. + zh_Hans: 通过其 ID 从指定的 Trello 看板检索所有列表,提供看板组织和当前阶段或类别的概览。 + pt_BR: Recupera todas as listas de um quadro Trello especificado pelo seu ID, fornecendo uma visão geral da organização do quadro e das fases ou categorias atuais. + llm: Fetch and display all lists from a specific Trello board by inputting the board's ID. This aids in understanding the board's structure and task categorization. +parameters: + - name: boardId + type: string + required: true + label: + en_US: Board ID + zh_Hans: 看板 ID + pt_BR: ID do Quadro + human_description: + en_US: The unique identifier of the Trello board from which to retrieve the lists. + zh_Hans: 用于检索列表的 Trello 看板的唯一标识符。 + pt_BR: O identificador único do quadro Trello do qual as listas serão recuperadas. + llm_description: Enter the ID of the Trello board to obtain a detailed list of all its lists, providing insight into the board's structure. + form: llm diff --git a/api/core/tools/provider/builtin/trello/tools/update_board.py b/api/core/tools/provider/builtin/trello/tools/update_board.py new file mode 100644 index 0000000000000000000000000000000000000000..1e358b00f49add31e10dc2ff6bf18eb236611986 --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/update_board.py @@ -0,0 +1,47 @@ +from typing import Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class UpdateBoardByIdTool(BuiltinTool): + """ + Tool for updating a Trello board by its ID with various parameters. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Union[str, int, bool, None]]) -> ToolInvokeMessage: + """ + Invoke the tool to update a Trello board by its ID. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (dict[str, Union[str, int, bool, None]]): The parameters for the tool invocation, + including board ID and updates. + + Returns: + ToolInvokeMessage: The result of the tool invocation. + """ + api_key = self.runtime.credentials.get("trello_api_key") + token = self.runtime.credentials.get("trello_api_token") + board_id = tool_parameters.pop("boardId", None) + + if not (api_key and token and board_id): + return self.create_text_message("Missing required parameters: API key, token, or board ID.") + + url = f"https://api.trello.com/1/boards/{board_id}" + + # Removing parameters not intended for update action or with None value + params = {k: v for k, v in tool_parameters.items() if v is not None} + params["key"] = api_key + params["token"] = token + + try: + response = requests.put(url, params=params) + response.raise_for_status() + except requests.exceptions.RequestException as e: + return self.create_text_message("Failed to update board") + + updated_board = response.json() + return self.create_text_message(text=f"Board '{updated_board['name']}' updated successfully.") diff --git a/api/core/tools/provider/builtin/trello/tools/update_board.yaml b/api/core/tools/provider/builtin/trello/tools/update_board.yaml new file mode 100644 index 0000000000000000000000000000000000000000..487919631ade3433719b671137b4260b4a9078cf --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/update_board.yaml @@ -0,0 +1,157 @@ +identity: + name: update_board_by_id + author: Yash Parmar + label: + en_US: Update Board by ID + zh_Hans: 通过 ID 更新看板 + pt_BR: Atualizar Quadro por ID +description: + human: + en_US: Updates a Trello board's settings based on the provided ID and parameters. Allows for changing the board's name, description, status, and other preferences. + zh_Hans: 根据提供的 ID 和参数更新 Trello 看板的设置。允许更改看板的名称、描述、状态和其他偏好设置。 + pt_BR: Atualiza as configurações de um quadro Trello com base no ID fornecido e nos parâmetros. Permite alterar o nome, descrição, status e outras preferências do quadro. + llm: Modify a Trello board's attributes like its name, description, and visibility settings using the board's ID. This tool streamlines board customization and management. +parameters: + - name: boardId + type: string + required: true + label: + en_US: Board ID + zh_Hans: 看板 ID + pt_BR: ID do Quadro + human_description: + en_US: The unique identifier of the Trello board you want to update. Ensures targeted and precise updates. + zh_Hans: 您要更新的 Trello 看板的唯一标识符。确保目标准确和更新精确。 + pt_BR: O identificador único do quadro Trello que você deseja atualizar. Garante atualizações direcionadas e precisas. + llm_description: Provide the specific ID of the Trello board you aim to update to ensure accuracy in modification process. + form: llm + - name: name + type: string + required: false + label: + en_US: Board Name + zh_Hans: 看板名称 + pt_BR: Nome do Quadro + human_description: + en_US: Optional. The new name for the board. + zh_Hans: 可选。看板的新名称。 + pt_BR: Opcional. O novo nome para o quadro. + llm_description: Enter a new name for the board if you wish to change it; this name identifies the board in Trello. + form: llm + - name: desc + type: string + required: false + label: + en_US: Board Description + zh_Hans: 看板描述 + pt_BR: Descrição do Quadro + human_description: + en_US: Optional. The new description for the board. + zh_Hans: 可选。看板的新描述。 + pt_BR: Opcional. A nova descrição para o quadro. + llm_description: Provide a new description for the board if you wish to update it; this description provides additional context about the board. + form: llm + - name: closed + type: boolean + required: false + label: + en_US: Closed + zh_Hans: 已关闭 + pt_BR: Fechado + human_description: + en_US: Optional. Set to true to close the board, or false to keep it open. + zh_Hans: 可选。设置为 true 以关闭看板,或设置为 false 以保持打开。 + pt_BR: Opcional. Defina como true para fechar o quadro ou como false para mantê-lo aberto. + llm_description: Specify whether the board should be closed or kept open by setting this parameter to true or false. + form: llm + - name: subscribed + type: string + required: false + label: + en_US: Subscribed + zh_Hans: 订阅 + pt_BR: Inscrito + human_description: + en_US: Optional. Set to true to subscribe to the board, or false to unsubscribe. + zh_Hans: 可选。设置为 true 以订阅看板,或设置为 false 以取消订阅。 + pt_BR: Opcional. Defina como true para se inscrever no quadro ou como false para cancelar a inscrição. + llm_description: Choose to subscribe or unsubscribe from the board by setting this parameter to true or false. + form: llm + - name: idOrganization + type: string + required: false + label: + en_US: Organization ID + zh_Hans: 组织 ID + pt_BR: ID da Organização + human_description: + en_US: Optional. The ID of the organization to which the board belongs. + zh_Hans: 可选。看板所属组织的 ID。 + pt_BR: Opcional. O ID da organização à qual o quadro pertence. + llm_description: Input the ID of the organization to which the board is associated, if applicable. + form: llm + - name: prefs_permissionLevel + type: string + required: false + label: + en_US: Permission Level + zh_Hans: 权限级别 + pt_BR: Nível de Permissão + human_description: + en_US: Optional. The permission level for the board. Valid values are 'private', 'org', or 'public'. + zh_Hans: 可选。看板的权限级别。有效值为 'private'、'org' 或 'public'。 + pt_BR: Opcional. O nível de permissão para o quadro. Os valores válidos são 'private', 'org' ou 'public'. + llm_description: Specify the permission level for the board by choosing from 'private', 'org', or 'public'. + form: llm + - name: prefs_selfJoin + type: boolean + required: false + label: + en_US: Allow Self-Join + zh_Hans: 允许自行加入 + pt_BR: Permitir Auto-Inscrição + human_description: + en_US: Optional. Set to true to allow members to join the board without an invitation, or false to require an invitation. + zh_Hans: 可选。设置为 true 以允许成员加入看板而无需邀请,或设置为 false 以要求邀请。 + pt_BR: Opcional. Defina como true para permitir que os membros se inscrevam no quadro sem um convite, ou como false para exigir um convite. + llm_description: Choose whether to allow members to join the board without an invitation by setting this parameter to true or false. + form: llm + - name: prefs_cardCovers + type: boolean + required: false + label: + en_US: Card Covers + zh_Hans: 卡片封面 + pt_BR: Capas de Cartão + human_description: + en_US: Optional. Set to true to enable card covers, or false to disable them. + zh_Hans: 可选。设置为 true 以启用卡片封面,或设置为 false 以禁用卡片封面。 + pt_BR: Opcional. Defina como true para habilitar capas de cartão ou como false para desabilitá-las. + llm_description: Enable or disable card covers by setting this parameter to true or false. + form: llm + - name: prefs_hideVotes + type: boolean + required: false + label: + en_US: Hide Votes + zh_Hans: 隐藏投票 + pt_BR: Ocultar Votos + human_description: + en_US: Optional. Set to true to hide votes, or false to show them. + zh_Hans: 可选。设置为 true 以隐藏投票,或设置为 false 以显示投票。 + pt_BR: Opcional. Defina como true para ocultar votos ou como false para mostrá-los. + llm_description: Choose to hide or show votes by setting this parameter to true or false. + form: llm + - name: prefs_invitations + type: string + required: false + label: + en_US: Invitations + zh_Hans: 邀请 + pt_BR: Convites + human_description: + en_US: Optional. Set to 'members' to allow only board members to send invitations, or 'admins' to allow admins to send invitations. + zh_Hans: 可选。设置为 'members' 以仅允许看板成员发送邀请,或设置为 'admins' 以允许管理员发送邀请。 + pt_BR: Opcional. Defina como 'members' para permitir que apenas membros do quadro enviem convites, ou 'admins' para permitir que os administradores enviem convites. + llm_description: Choose who can send invitations by setting this parameter to 'members' or 'admins'. + form: llm diff --git a/api/core/tools/provider/builtin/trello/tools/update_card.py b/api/core/tools/provider/builtin/trello/tools/update_card.py new file mode 100644 index 0000000000000000000000000000000000000000..d25fcbafaa6326695ab9f4ab2ddeeeab1b355a06 --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/update_card.py @@ -0,0 +1,45 @@ +from typing import Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class UpdateCardByIdTool(BuiltinTool): + """ + Tool for updating a Trello card by its ID. + """ + + def _invoke(self, user_id: str, tool_parameters: dict[str, Union[str, int, bool, None]]) -> ToolInvokeMessage: + """ + Invoke the tool to update a Trello card by its ID. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (dict[str, Union[str, int, bool, None]]): The parameters for the tool invocation, + including the card ID and updates. + + Returns: + ToolInvokeMessage: The result of the tool invocation. + """ + api_key = self.runtime.credentials.get("trello_api_key") + token = self.runtime.credentials.get("trello_api_token") + card_id = tool_parameters.get("id") + + if not (api_key and token and card_id): + return self.create_text_message("Missing required parameters: API key, token, or card ID.") + + # Constructing the URL and the payload for the PUT request + url = f"https://api.trello.com/1/cards/{card_id}" + params = {k: v for k, v in tool_parameters.items() if v is not None and k != "id"} + params.update({"key": api_key, "token": token}) + + try: + response = requests.put(url, params=params) + response.raise_for_status() + except requests.exceptions.RequestException as e: + return self.create_text_message("Failed to update card") + + updated_card_info = f"Card '{card_id}' updated successfully." + return self.create_text_message(text=updated_card_info) diff --git a/api/core/tools/provider/builtin/trello/tools/update_card.yaml b/api/core/tools/provider/builtin/trello/tools/update_card.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5240dfc3ed25647d36af5b55ddceb314b61bc942 --- /dev/null +++ b/api/core/tools/provider/builtin/trello/tools/update_card.yaml @@ -0,0 +1,81 @@ +identity: + name: update_card_by_id + author: Yash Parmar + label: + en_US: Update Card by ID + zh_Hans: 通过 ID 更新卡片 + pt_BR: Atualizar Cartão por ID +description: + human: + en_US: Updates specified attributes of a Trello card, such as its name, description, list ID, and board ID, by providing the card's unique ID. + zh_Hans: 通过提供卡片的唯一 ID,更新 Trello 卡片的特定属性,如其名称、描述、列表 ID 和看板 ID。 + pt_BR: Atualiza atributos específicos de um cartão Trello, como seu nome, descrição, ID da lista e ID do quadro, fornecendo o ID único do cartão. + llm: Modify a Trello card's key details, including name, description, and its placement on the board, by using the card's ID. Enables precise and targeted updates to card information. +parameters: + - name: id + type: string + required: true + label: + en_US: Card ID + zh_Hans: 卡片 ID + pt_BR: ID do Cartão + human_description: + en_US: The unique identifier of the Trello card you intend to update. + zh_Hans: 您打算更新的 Trello 卡片的唯一标识符。 + pt_BR: O identificador único do cartão Trello que você pretende atualizar. + llm_description: Input the ID of the Trello card to be updated to ensure the correct card is targeted. + form: llm + # Include other parameters following the same pattern + - name: name + type: string + required: false + label: + en_US: New Name + zh_Hans: 新名称 + pt_BR: Novo Nome + human_description: + en_US: Optional. The new name to assign to the card. + zh_Hans: 可选。要分配给卡片的新名称。 + pt_BR: Opcional. O novo nome a ser atribuído ao cartão. + llm_description: Specify a new name for the card if changing it. This name is what will be displayed on the Trello board. + form: llm + # Add definitions for desc, idList and idBoard parameters + - name: desc + type: string + required: false + label: + en_US: New Description + zh_Hans: 新描述 + pt_BR: Nova Descrição + human_description: + en_US: Optional. The new description to assign to the card. + zh_Hans: 可选。要分配给卡片的新描述。 + pt_BR: Opcional. A nova descrição a ser atribuída ao cartão. + llm_description: Provide a new description for the card if you wish to update it; this description provides additional context about the card. + form: llm + - name: idList + type: string + required: false + label: + en_US: List ID + zh_Hans: 列表 ID + pt_BR: ID da Lista + human_description: + en_US: Optional. The ID of the list to which the card should be moved. + zh_Hans: 可选。卡片应移动到的列表的 ID。 + pt_BR: Opcional. O ID da lista para a qual o cartão deve ser movido. + llm_description: Enter the ID of the list where you want to move the card. This action relocates the card to the specified list. + form: llm + - name: idBoard + type: string + required: false + label: + en_US: Board ID + zh_Hans: 看板 ID + pt_BR: ID do Quadro + human_description: + en_US: Optional. The ID of the board to which the card should be moved. + zh_Hans: 可选。卡片应移动到的看板的 ID。 + pt_BR: Opcional. O ID do quadro para o qual o cartão deve ser movido. + llm_description: Provide the ID of the board where you want to move the card. This action relocates the card to the specified board. + form: llm diff --git a/api/core/tools/provider/builtin/trello/trello.py b/api/core/tools/provider/builtin/trello/trello.py new file mode 100644 index 0000000000000000000000000000000000000000..e0dca50ec99aeed1ab76dc4581f37e37d76058eb --- /dev/null +++ b/api/core/tools/provider/builtin/trello/trello.py @@ -0,0 +1,34 @@ +from typing import Any + +import requests + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class TrelloProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + """Validate Trello API credentials by making a test API call. + + Args: + credentials (dict[str, Any]): The Trello API credentials to validate. + + Raises: + ToolProviderCredentialValidationError: If the credentials are invalid. + """ + api_key = credentials.get("trello_api_key") + token = credentials.get("trello_api_token") + url = f"https://api.trello.com/1/members/me?key={api_key}&token={token}" + + try: + response = requests.get(url) + response.raise_for_status() # Raises an HTTPError for bad responses + except requests.exceptions.HTTPError as e: + if response.status_code == 401: + # Unauthorized, indicating invalid credentials + raise ToolProviderCredentialValidationError("Invalid Trello credentials: Unauthorized.") + # Handle other potential HTTP errors + raise ToolProviderCredentialValidationError("Error validating Trello credentials") + except requests.exceptions.RequestException as e: + # Handle other exceptions, such as connection errors + raise ToolProviderCredentialValidationError("Error validating Trello credentials") diff --git a/api/core/tools/provider/builtin/trello/trello.yaml b/api/core/tools/provider/builtin/trello/trello.yaml new file mode 100644 index 0000000000000000000000000000000000000000..49c9f4f9a178f8348f7b35eb8773c86c4ffa06aa --- /dev/null +++ b/api/core/tools/provider/builtin/trello/trello.yaml @@ -0,0 +1,47 @@ +identity: + author: Yash Parmar + name: trello + label: + en_US: Trello + zh_Hans: Trello + pt_BR: Trello + description: + en_US: "Trello: A visual tool for organizing your work and life." + zh_Hans: "Trello: 一个用于组织工作和生活的视觉工具。" + pt_BR: "Trello: Uma ferramenta visual para organizar seu trabalho e vida." + icon: icon.svg + tags: + - productivity +credentials_for_provider: + trello_api_key: + type: secret-input + required: true + label: + en_US: Trello API key + zh_Hans: Trello API key + pt_BR: Trello API key + placeholder: + en_US: Enter your Trello API key + zh_Hans: 输入您的 Trello API key + pt_BR: Insira sua chave API do Trello + help: + en_US: Obtain your API key from Trello's website. + zh_Hans: 从 Trello 网站获取您的 API key。 + pt_BR: Obtenha sua chave API no site do Trello. + url: https://developer.atlassian.com/cloud/trello/guides/rest-api/api-introduction/ + trello_api_token: + type: secret-input + required: true + label: + en_US: Trello API token + zh_Hans: Trello API token + pt_BR: Trello API token + placeholder: + en_US: Enter your Trello API token + zh_Hans: 输入您的 Trello API token + pt_BR: Insira seu token API do Trello + help: + en_US: Secure your API token from Trello's website. + zh_Hans: 从 Trello 网站获取您的 API token。 + pt_BR: Garanta seu token API no site do Trello. + url: https://developer.atlassian.com/cloud/trello/guides/rest-api/api-introduction/ diff --git a/api/core/tools/provider/builtin/twilio/_assets/icon.svg b/api/core/tools/provider/builtin/twilio/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..a1e2bd12c27d64dd9811534608a47e9b36e3e74c --- /dev/null +++ b/api/core/tools/provider/builtin/twilio/_assets/icon.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/twilio/tools/send_message.py b/api/core/tools/provider/builtin/twilio/tools/send_message.py new file mode 100644 index 0000000000000000000000000000000000000000..98a108f4ec7e93e299a44b3ff737abc9f215a8ea --- /dev/null +++ b/api/core/tools/provider/builtin/twilio/tools/send_message.py @@ -0,0 +1,97 @@ +from typing import Any, Optional, Union + +from pydantic import BaseModel, field_validator + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class TwilioAPIWrapper(BaseModel): + """Messaging Client using Twilio. + + To use, you should have the ``twilio`` python package installed, + and the environment variables ``TWILIO_ACCOUNT_SID``, ``TWILIO_AUTH_TOKEN``, and + ``TWILIO_FROM_NUMBER``, or pass `account_sid`, `auth_token`, and `from_number` as + named parameters to the constructor. + """ + + client: Any = None #: :meta private: + account_sid: Optional[str] = None + """Twilio account string identifier.""" + auth_token: Optional[str] = None + """Twilio auth token.""" + from_number: Optional[str] = None + """A Twilio phone number in [E.164](https://www.twilio.com/docs/glossary/what-e164) + format, an + [alphanumeric sender ID](https://www.twilio.com/docs/sms/send-messages#use-an-alphanumeric-sender-id), + or a [Channel Endpoint address](https://www.twilio.com/docs/sms/channels#channel-addresses) + that is enabled for the type of message you want to send. Phone numbers or + [short codes](https://www.twilio.com/docs/sms/api/short-code) purchased from + Twilio also work here. You cannot, for example, spoof messages from a private + cell phone number. If you are using `messaging_service_sid`, this parameter + must be empty. + """ + + @field_validator("client", mode="before") + @classmethod + def set_validator(cls, values: dict) -> dict: + """Validate that api key and python package exists in environment.""" + try: + from twilio.rest import Client # type: ignore + except ImportError: + raise ImportError("Could not import twilio python package. Please install it with `pip install twilio`.") + account_sid = values.get("account_sid") + auth_token = values.get("auth_token") + values["from_number"] = values.get("from_number") + values["client"] = Client(account_sid, auth_token) + + return values + + def run(self, body: str, to: str) -> str: + """Run body through Twilio and respond with message sid. + + Args: + body: The text of the message you want to send. Can be up to 1,600 + characters in length. + to: The destination phone number in + [E.164](https://www.twilio.com/docs/glossary/what-e164) format for + SMS/MMS or + [Channel user address](https://www.twilio.com/docs/sms/channels#channel-addresses) + for other 3rd-party channels. + """ + message = self.client.messages.create(to, from_=self.from_number, body=body) + return message.sid + + +class SendMessageTool(BuiltinTool): + """ + A tool for sending messages using Twilio API. + + Args: + user_id (str): The ID of the user invoking the tool. + tool_parameters (Dict[str, Any]): The parameters required for sending the message. + + Returns: + Union[ToolInvokeMessage, List[ToolInvokeMessage]]: The result of invoking the tool, + which includes the status of the message sending operation. + """ + + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + account_sid = self.runtime.credentials["account_sid"] + auth_token = self.runtime.credentials["auth_token"] + from_number = self.runtime.credentials["from_number"] + + message = tool_parameters["message"] + to_number = tool_parameters["to_number"] + + if to_number.startswith("whatsapp:"): + from_number = f"whatsapp: {from_number}" + + twilio = TwilioAPIWrapper(account_sid=account_sid, auth_token=auth_token, from_number=from_number) + + # Sending the message through Twilio + result = twilio.run(message, to_number) + + return self.create_text_message(text="Message sent successfully.") diff --git a/api/core/tools/provider/builtin/twilio/tools/send_message.yaml b/api/core/tools/provider/builtin/twilio/tools/send_message.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e129698c86aeb60bb77fa9dff8dd71e9476697b3 --- /dev/null +++ b/api/core/tools/provider/builtin/twilio/tools/send_message.yaml @@ -0,0 +1,40 @@ +identity: + name: send_message + author: Yash Parmar + label: + en_US: SendMessage + zh_Hans: 发送消息 + pt_BR: SendMessage +description: + human: + en_US: Send SMS or Twilio Messaging Channels messages. + zh_Hans: 发送SMS或Twilio消息通道消息。 + pt_BR: Send SMS or Twilio Messaging Channels messages. + llm: Send SMS or Twilio Messaging Channels messages. Supports different channels including WhatsApp. +parameters: + - name: message + type: string + required: true + label: + en_US: Message + zh_Hans: 消息内容 + pt_BR: Message + human_description: + en_US: The content of the message to be sent. + zh_Hans: 要发送的消息内容。 + pt_BR: The content of the message to be sent. + llm_description: The content of the message to be sent. + form: llm + - name: to_number + type: string + required: true + label: + en_US: To Number + zh_Hans: 收信号码 + pt_BR: Para Número + human_description: + en_US: The recipient's phone number. Prefix with 'whatsapp:' for WhatsApp messages, e.g., "whatsapp:+1234567890". + zh_Hans: 收件人的电话号码。WhatsApp消息前缀为'whatsapp:',例如,"whatsapp:+1234567890"。 + pt_BR: The recipient's phone number. Prefix with 'whatsapp:' for WhatsApp messages, e.g., "whatsapp:+1234567890". + llm_description: The recipient's phone number. Prefix with 'whatsapp:' for WhatsApp messages, e.g., "whatsapp:+1234567890". + form: llm diff --git a/api/core/tools/provider/builtin/twilio/twilio.py b/api/core/tools/provider/builtin/twilio/twilio.py new file mode 100644 index 0000000000000000000000000000000000000000..649e03d185121c5c78ebf004d088ea68e0cd0be6 --- /dev/null +++ b/api/core/tools/provider/builtin/twilio/twilio.py @@ -0,0 +1,29 @@ +from typing import Any + +from twilio.base.exceptions import TwilioRestException # type: ignore +from twilio.rest import Client # type: ignore + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class TwilioProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + # Extract credentials + account_sid = credentials["account_sid"] + auth_token = credentials["auth_token"] + from_number = credentials["from_number"] + + # Initialize twilio client + client = Client(account_sid, auth_token) + + # fetch account + client.api.accounts(account_sid).fetch() + + except TwilioRestException as e: + raise ToolProviderCredentialValidationError(f"Twilio API error: {e.msg}") from e + except KeyError as e: + raise ToolProviderCredentialValidationError(f"Missing required credential: {e}") from e + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/twilio/twilio.yaml b/api/core/tools/provider/builtin/twilio/twilio.yaml new file mode 100644 index 0000000000000000000000000000000000000000..21867c1da5dc324d263b05aebc7ecd74c3a5a47b --- /dev/null +++ b/api/core/tools/provider/builtin/twilio/twilio.yaml @@ -0,0 +1,48 @@ +identity: + author: Yash Parmar + name: twilio + label: + en_US: Twilio + zh_Hans: Twilio + pt_BR: Twilio + description: + en_US: Send messages through SMS or Twilio Messaging Channels. + zh_Hans: 通过SMS或Twilio消息通道发送消息。 + pt_BR: Send messages through SMS or Twilio Messaging Channels. + icon: icon.svg + tags: + - social +credentials_for_provider: + account_sid: + type: secret-input + required: true + label: + en_US: Account SID + zh_Hans: 账户SID + pt_BR: Account SID + placeholder: + en_US: Please input your Twilio Account SID + zh_Hans: 请输入您的Twilio账户SID + pt_BR: Please input your Twilio Account SID + auth_token: + type: secret-input + required: true + label: + en_US: Auth Token + zh_Hans: 认证令牌 + pt_BR: Auth Token + placeholder: + en_US: Please input your Twilio Auth Token + zh_Hans: 请输入您的Twilio认证令牌 + pt_BR: Please input your Twilio Auth Token + from_number: + type: secret-input + required: true + label: + en_US: From Number + zh_Hans: 发信号码 + pt_BR: De Número + placeholder: + en_US: Please input your Twilio phone number + zh_Hans: 请输入您的Twilio电话号码 + pt_BR: Please input your Twilio phone number diff --git a/api/core/tools/provider/builtin/vanna/_assets/icon.png b/api/core/tools/provider/builtin/vanna/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..3a9011b54d8a07f01e6b2fb934f3937bca0fd85a Binary files /dev/null and b/api/core/tools/provider/builtin/vanna/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/vanna/tools/vanna.py b/api/core/tools/provider/builtin/vanna/tools/vanna.py new file mode 100644 index 0000000000000000000000000000000000000000..a6afd2dddfc63ae1dba7893541e3db9e4a6fba98 --- /dev/null +++ b/api/core/tools/provider/builtin/vanna/tools/vanna.py @@ -0,0 +1,134 @@ +from typing import Any, Union + +from vanna.remote import VannaDefault # type: ignore + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.tool.builtin_tool import BuiltinTool + + +class VannaTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + # Ensure runtime and credentials + if not self.runtime or not self.runtime.credentials: + raise ToolProviderCredentialValidationError("Tool runtime or credentials are missing") + api_key = self.runtime.credentials.get("api_key", None) + if not api_key: + raise ToolProviderCredentialValidationError("Please input api key") + + model = tool_parameters.get("model", "") + if not model: + return self.create_text_message("Please input RAG model") + + prompt = tool_parameters.get("prompt", "") + if not prompt: + return self.create_text_message("Please input prompt") + + url = tool_parameters.get("url", "") + if not url: + return self.create_text_message("Please input URL/Host/DSN") + + db_name = tool_parameters.get("db_name", "") + username = tool_parameters.get("username", "") + password = tool_parameters.get("password", "") + port = tool_parameters.get("port", 0) + + base_url = self.runtime.credentials.get("base_url", None) + vn = VannaDefault(model=model, api_key=api_key, config={"endpoint": base_url}) + + db_type = tool_parameters.get("db_type", "") + if db_type in {"Postgres", "MySQL", "Hive", "ClickHouse"}: + if not db_name: + return self.create_text_message("Please input database name") + if not username: + return self.create_text_message("Please input username") + if port < 1: + return self.create_text_message("Please input port") + + schema_sql = "SELECT * FROM INFORMATION_SCHEMA.COLUMNS" + match db_type: + case "SQLite": + schema_sql = "SELECT type, sql FROM sqlite_master WHERE sql is not null" + vn.connect_to_sqlite(url) + case "Postgres": + vn.connect_to_postgres(host=url, dbname=db_name, user=username, password=password, port=port) + case "DuckDB": + vn.connect_to_duckdb(url=url) + case "SQLServer": + vn.connect_to_mssql(url) + case "MySQL": + vn.connect_to_mysql(host=url, dbname=db_name, user=username, password=password, port=port) + case "Oracle": + vn.connect_to_oracle(user=username, password=password, dsn=url) + case "Hive": + vn.connect_to_hive(host=url, dbname=db_name, user=username, password=password, port=port) + case "ClickHouse": + vn.connect_to_clickhouse(host=url, dbname=db_name, user=username, password=password, port=port) + + enable_training = tool_parameters.get("enable_training", False) + reset_training_data = tool_parameters.get("reset_training_data", False) + if enable_training: + if reset_training_data: + existing_training_data = vn.get_training_data() + if len(existing_training_data) > 0: + for _, training_data in existing_training_data.iterrows(): + vn.remove_training_data(training_data["id"]) + + ddl = tool_parameters.get("ddl", "") + question = tool_parameters.get("question", "") + sql = tool_parameters.get("sql", "") + memos = tool_parameters.get("memos", "") + training_metadata = tool_parameters.get("training_metadata", False) + + if training_metadata: + if db_type == "SQLite": + df_ddl = vn.run_sql(schema_sql) + for ddl in df_ddl["sql"].to_list(): + vn.train(ddl=ddl) + else: + df_information_schema = vn.run_sql(schema_sql) + plan = vn.get_training_plan_generic(df_information_schema) + vn.train(plan=plan) + + if ddl: + vn.train(ddl=ddl) + + if sql: + if question: + vn.train(question=question, sql=sql) + else: + vn.train(sql=sql) + if memos: + vn.train(documentation=memos) + + ######################################################################################### + # Due to CVE-2024-5565, we have to disable the chart generation feature + # The Vanna library uses a prompt function to present the user with visualized results, + # it is possible to alter the prompt using prompt injection and run arbitrary Python code + # instead of the intended visualization code. + # Specifically - allowing external input to the library’s “ask” method + # with "visualize" set to True (default behavior) leads to remote code execution. + # Affected versions: <= 0.5.5 + ######################################################################################### + allow_llm_to_see_data = tool_parameters.get("allow_llm_to_see_data", False) + res = vn.ask( + prompt, print_results=False, auto_train=True, visualize=False, allow_llm_to_see_data=allow_llm_to_see_data + ) + + result = [] + + if res is not None: + result.append(self.create_text_message(res[0])) + if len(res) > 1 and res[1] is not None: + result.append(self.create_text_message(res[1].to_markdown())) + if len(res) > 2 and res[2] is not None: + result.append( + self.create_blob_message(blob=res[2].to_image(format="svg"), meta={"mime_type": "image/svg+xml"}) + ) + + return result diff --git a/api/core/tools/provider/builtin/vanna/tools/vanna.yaml b/api/core/tools/provider/builtin/vanna/tools/vanna.yaml new file mode 100644 index 0000000000000000000000000000000000000000..309681321b1f3f0e1062065bf3b51b1bcd604c81 --- /dev/null +++ b/api/core/tools/provider/builtin/vanna/tools/vanna.yaml @@ -0,0 +1,213 @@ +identity: + name: vanna + author: QCTC + label: + en_US: Vanna.AI + zh_Hans: Vanna.AI +description: + human: + en_US: The fastest way to get actionable insights from your database just by asking questions. + zh_Hans: 一个基于大模型和RAG的Text2SQL工具。 + llm: A tool for converting text to SQL. +parameters: + - name: prompt + type: string + required: true + label: + en_US: Prompt + zh_Hans: 提示词 + pt_BR: Prompt + human_description: + en_US: used for generating SQL + zh_Hans: 用于生成SQL + llm_description: key words for generating SQL + form: llm + - name: model + type: string + required: true + label: + en_US: RAG Model + zh_Hans: RAG模型 + human_description: + en_US: RAG Model for your database DDL + zh_Hans: 存储数据库训练数据的RAG模型 + llm_description: RAG Model for generating SQL + form: llm + - name: db_type + type: select + required: true + options: + - value: SQLite + label: + en_US: SQLite + zh_Hans: SQLite + - value: Postgres + label: + en_US: Postgres + zh_Hans: Postgres + - value: DuckDB + label: + en_US: DuckDB + zh_Hans: DuckDB + - value: SQLServer + label: + en_US: Microsoft SQL Server + zh_Hans: 微软 SQL Server + - value: MySQL + label: + en_US: MySQL + zh_Hans: MySQL + - value: Oracle + label: + en_US: Oracle + zh_Hans: Oracle + - value: Hive + label: + en_US: Hive + zh_Hans: Hive + - value: ClickHouse + label: + en_US: ClickHouse + zh_Hans: ClickHouse + default: SQLite + label: + en_US: DB Type + zh_Hans: 数据库类型 + human_description: + en_US: Database type. + zh_Hans: 选择要链接的数据库类型。 + form: form + - name: url + type: string + required: true + label: + en_US: URL/Host/DSN + zh_Hans: URL/Host/DSN + human_description: + en_US: Please input depending on DB type, visit https://vanna.ai/docs/ for more specification + zh_Hans: 请根据数据库类型,填入对应值,详情参考https://vanna.ai/docs/ + form: form + - name: db_name + type: string + required: false + label: + en_US: DB name + zh_Hans: 数据库名 + human_description: + en_US: Database name + zh_Hans: 数据库名 + form: form + - name: username + type: string + required: false + label: + en_US: Username + zh_Hans: 用户名 + human_description: + en_US: Username + zh_Hans: 用户名 + form: form + - name: password + type: secret-input + required: false + label: + en_US: Password + zh_Hans: 密码 + human_description: + en_US: Password + zh_Hans: 密码 + form: form + - name: port + type: number + required: false + label: + en_US: Port + zh_Hans: 端口 + human_description: + en_US: Port + zh_Hans: 端口 + form: form + - name: ddl + type: string + required: false + label: + en_US: Training DDL + zh_Hans: 训练DDL + human_description: + en_US: DDL statements for training data + zh_Hans: 用于训练RAG Model的建表语句 + form: llm + - name: question + type: string + required: false + label: + en_US: Training Question + zh_Hans: 训练问题 + human_description: + en_US: Question-SQL Pairs + zh_Hans: Question-SQL中的问题 + form: llm + - name: sql + type: string + required: false + label: + en_US: Training SQL + zh_Hans: 训练SQL + human_description: + en_US: SQL queries to your training data + zh_Hans: 用于训练RAG Model的SQL语句 + form: llm + - name: memos + type: string + required: false + label: + en_US: Training Memos + zh_Hans: 训练说明 + human_description: + en_US: Sometimes you may want to add documentation about your business terminology or definitions + zh_Hans: 添加更多关于数据库的业务说明 + form: llm + - name: enable_training + type: boolean + required: false + default: false + label: + en_US: Training Data + zh_Hans: 训练数据 + human_description: + en_US: You only need to train once. Do not train again unless you want to add more training data + zh_Hans: 训练数据无更新时,训练一次即可 + form: form + - name: reset_training_data + type: boolean + required: false + default: false + label: + en_US: Reset Training Data + zh_Hans: 重置训练数据 + human_description: + en_US: Remove all training data in the current RAG Model + zh_Hans: 删除当前RAG Model中的所有训练数据 + form: form + - name: training_metadata + type: boolean + required: false + default: false + label: + en_US: Training Metadata + zh_Hans: 训练元数据 + human_description: + en_US: If enabled, it will attempt to train on the metadata of that database + zh_Hans: 是否自动从数据库获取元数据来训练 + form: form + - name: allow_llm_to_see_data + type: boolean + required: false + default: false + label: + en_US: Whether to allow the LLM to see the data + zh_Hans: 是否允许LLM查看数据 + human_description: + en_US: Whether to allow the LLM to see the data + zh_Hans: 是否允许LLM查看数据 + form: form diff --git a/api/core/tools/provider/builtin/vanna/vanna.py b/api/core/tools/provider/builtin/vanna/vanna.py new file mode 100644 index 0000000000000000000000000000000000000000..4f9cac2beb01bb24c7be6fcdba84bb9d167f1773 --- /dev/null +++ b/api/core/tools/provider/builtin/vanna/vanna.py @@ -0,0 +1,46 @@ +import re +from typing import Any +from urllib.parse import urlparse + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.vanna.tools.vanna import VannaTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class VannaProvider(BuiltinToolProviderController): + def _get_protocol_and_main_domain(self, url): + parsed_url = urlparse(url) + protocol = parsed_url.scheme + hostname = parsed_url.hostname + port = f":{parsed_url.port}" if parsed_url.port else "" + + # Check if the hostname is an IP address + is_ip = re.match(r"^\d{1,3}(\.\d{1,3}){3}$", hostname) is not None + + # Return the full hostname (with port if present) for IP addresses, otherwise return the main domain + main_domain = f"{hostname}{port}" if is_ip else ".".join(hostname.split(".")[-2:]) + port + return f"{protocol}://{main_domain}" + + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + base_url = credentials.get("base_url") + if not base_url: + base_url = "https://ask.vanna.ai/rpc" + else: + base_url = base_url.removesuffix("/") + credentials["base_url"] = base_url + try: + VannaTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "model": "chinook", + "db_type": "SQLite", + "url": f"{self._get_protocol_and_main_domain(credentials['base_url'])}/Chinook.sqlite", + "query": "What are the top 10 customers by sales?", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/vanna/vanna.yaml b/api/core/tools/provider/builtin/vanna/vanna.yaml new file mode 100644 index 0000000000000000000000000000000000000000..cf3fdca562c0b3d151353af15fe92d3a6d9f1bc5 --- /dev/null +++ b/api/core/tools/provider/builtin/vanna/vanna.yaml @@ -0,0 +1,35 @@ +identity: + author: QCTC + name: vanna + label: + en_US: Vanna.AI + zh_Hans: Vanna.AI + description: + en_US: The fastest way to get actionable insights from your database just by asking questions. + zh_Hans: 一个基于大模型和RAG的Text2SQL工具。 + icon: icon.png + tags: + - utilities + - productivity +credentials_for_provider: + api_key: + type: secret-input + required: true + label: + en_US: API key + zh_Hans: API key + placeholder: + en_US: Please input your API key + zh_Hans: 请输入你的 API key + pt_BR: Please input your API key + help: + en_US: Get your API key from Vanna.AI + zh_Hans: 从 Vanna.AI 获取你的 API key + url: https://vanna.ai/account/profile + base_url: + type: text-input + required: false + label: + en_US: Vanna.AI Endpoint Base URL + placeholder: + en_US: https://ask.vanna.ai/rpc diff --git a/api/core/tools/provider/builtin/vectorizer/_assets/icon.png b/api/core/tools/provider/builtin/vectorizer/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..52f18db84372dcfc2968be27d75aec0fca430d55 Binary files /dev/null and b/api/core/tools/provider/builtin/vectorizer/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/vectorizer/tools/vectorizer.py b/api/core/tools/provider/builtin/vectorizer/tools/vectorizer.py new file mode 100644 index 0000000000000000000000000000000000000000..c722cd36c84e1593f30d5214b88107cb50208138 --- /dev/null +++ b/api/core/tools/provider/builtin/vectorizer/tools/vectorizer.py @@ -0,0 +1,82 @@ +from typing import Any, Union + +from httpx import post + +from core.file.enums import FileType +from core.file.file_manager import download +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter +from core.tools.errors import ToolParameterValidationError +from core.tools.tool.builtin_tool import BuiltinTool + + +class VectorizerTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + api_key_name = self.runtime.credentials.get("api_key_name") + api_key_value = self.runtime.credentials.get("api_key_value") + mode = tool_parameters.get("mode", "test") + + # image file for workflow mode + image = tool_parameters.get("image") + if image and image.type != FileType.IMAGE: + raise ToolParameterValidationError("Not a valid image") + # image_id for agent mode + image_id = tool_parameters.get("image_id", "") + + if image_id: + image_binary = self.get_variable_file(self.VariableKey.IMAGE) + if not image_binary: + return self.create_text_message("Image not found, please request user to generate image firstly.") + elif image: + image_binary = download(image) + else: + raise ToolParameterValidationError("Please provide either image or image_id") + + response = post( + "https://vectorizer.ai/api/v1/vectorize", + data={"mode": mode}, + files={"image": image_binary}, + auth=(api_key_name, api_key_value), + timeout=30, + ) + + if response.status_code != 200: + raise Exception(response.text) + + return [ + self.create_text_message("the vectorized svg is saved as an image."), + self.create_blob_message(blob=response.content, meta={"mime_type": "image/svg+xml"}), + ] + + def get_runtime_parameters(self) -> list[ToolParameter]: + """ + override the runtime parameters + """ + return [ + ToolParameter.get_simple_instance( + name="image_id", + llm_description=f"the image_id that you want to vectorize, \ + and the image_id should be specified in \ + {[i.name for i in self.list_default_image_variables()]}", + type=ToolParameter.ToolParameterType.SELECT, + required=False, + options=[i.name for i in self.list_default_image_variables()], + ), + ToolParameter( + name="image", + label=I18nObject(en_US="image", zh_Hans="image"), + human_description=I18nObject( + en_US="The image to be converted.", + zh_Hans="要转换的图片。", + ), + type=ToolParameter.ToolParameterType.FILE, + form=ToolParameter.ToolParameterForm.LLM, + llm_description="you should not input this parameter. just input the image_id.", + required=False, + ), + ] diff --git a/api/core/tools/provider/builtin/vectorizer/tools/vectorizer.yaml b/api/core/tools/provider/builtin/vectorizer/tools/vectorizer.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0afd1c201f9126c9fc09e12d30f4d689e528f0f0 --- /dev/null +++ b/api/core/tools/provider/builtin/vectorizer/tools/vectorizer.yaml @@ -0,0 +1,41 @@ +identity: + name: vectorizer + author: Dify + label: + en_US: Vectorizer.AI + zh_Hans: Vectorizer.AI +description: + human: + en_US: Convert your PNG and JPG images to SVG vectors quickly and easily. Fully automatically. Using AI. + zh_Hans: 一个将 PNG 和 JPG 图像快速轻松地转换为 SVG 矢量图的工具。 + llm: A tool for converting images to SVG vectors. you should input the image id as the input of this tool. the image id can be got from parameters. +parameters: + - name: image + type: file + label: + en_US: image + human_description: + en_US: The image to be converted. + zh_Hans: 要转换的图片。 + llm_description: you should not input this parameter. just input the image_id. + form: llm + - name: mode + type: select + required: true + options: + - value: production + label: + en_US: production + zh_Hans: 生产模式 + - value: test + label: + en_US: test + zh_Hans: 测试模式 + default: test + label: + en_US: Mode + zh_Hans: 模式 + human_description: + en_US: It is free to integrate with and test out the API in test mode, no subscription required. + zh_Hans: 在测试模式下,可以免费测试API。 + form: form diff --git a/api/core/tools/provider/builtin/vectorizer/vectorizer.py b/api/core/tools/provider/builtin/vectorizer/vectorizer.py new file mode 100644 index 0000000000000000000000000000000000000000..9d7613f8eaf1701c834c831f5cc4c9dd0d9331b5 --- /dev/null +++ b/api/core/tools/provider/builtin/vectorizer/vectorizer.py @@ -0,0 +1,8 @@ +from typing import Any + +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class VectorizerProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + return diff --git a/api/core/tools/provider/builtin/vectorizer/vectorizer.yaml b/api/core/tools/provider/builtin/vectorizer/vectorizer.yaml new file mode 100644 index 0000000000000000000000000000000000000000..94dae2087609d45ee92d79ac248ea23ed57a8ff1 --- /dev/null +++ b/api/core/tools/provider/builtin/vectorizer/vectorizer.yaml @@ -0,0 +1,39 @@ +identity: + author: Dify + name: vectorizer + label: + en_US: Vectorizer.AI + zh_Hans: Vectorizer.AI + description: + en_US: Convert your PNG and JPG images to SVG vectors quickly and easily. Fully automatically. Using AI. + zh_Hans: 一个将 PNG 和 JPG 图像快速轻松地转换为 SVG 矢量图的工具。 + icon: icon.png + tags: + - productivity + - image +credentials_for_provider: + api_key_name: + type: secret-input + required: true + label: + en_US: Vectorizer.AI API Key name + zh_Hans: Vectorizer.AI API Key name + placeholder: + en_US: Please input your Vectorizer.AI ApiKey name + zh_Hans: 请输入你的 Vectorizer.AI ApiKey name + help: + en_US: Get your Vectorizer.AI API Key from Vectorizer.AI. + zh_Hans: 从 Vectorizer.AI 获取您的 Vectorizer.AI API Key。 + url: https://vectorizer.ai/api + api_key_value: + type: secret-input + required: true + label: + en_US: Vectorizer.AI API Key + zh_Hans: Vectorizer.AI API Key + placeholder: + en_US: Please input your Vectorizer.AI ApiKey + zh_Hans: 请输入你的 Vectorizer.AI ApiKey + help: + en_US: Get your Vectorizer.AI API Key from Vectorizer.AI. + zh_Hans: 从 Vectorizer.AI 获取您的 Vectorizer.AI API Key。 diff --git a/api/core/tools/provider/builtin/webscraper/_assets/icon.svg b/api/core/tools/provider/builtin/webscraper/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..8123199a38a5e7b7218757fe7d8c83f3c86e0a1a --- /dev/null +++ b/api/core/tools/provider/builtin/webscraper/_assets/icon.svg @@ -0,0 +1,3 @@ + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/webscraper/tools/webscraper.py b/api/core/tools/provider/builtin/webscraper/tools/webscraper.py new file mode 100644 index 0000000000000000000000000000000000000000..12670b4b8b928939ec9aacb588203b64de4d346c --- /dev/null +++ b/api/core/tools/provider/builtin/webscraper/tools/webscraper.py @@ -0,0 +1,33 @@ +from typing import Any, Union + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.errors import ToolInvokeError +from core.tools.tool.builtin_tool import BuiltinTool + + +class WebscraperTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + try: + url = tool_parameters.get("url", "") + user_agent = tool_parameters.get("user_agent", "") + if not url: + return self.create_text_message("Please input url") + + # get webpage + result = self.get_url(url, user_agent=user_agent) + + if tool_parameters.get("generate_summary"): + # summarize and return + return self.create_text_message(self.summary(user_id=user_id, content=result)) + else: + # return full webpage + return self.create_text_message(result) + except Exception as e: + raise ToolInvokeError(str(e)) diff --git a/api/core/tools/provider/builtin/webscraper/tools/webscraper.yaml b/api/core/tools/provider/builtin/webscraper/tools/webscraper.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0bb48a941dcffe7a3c8ffdc2f9fbcb0cdf6590e1 --- /dev/null +++ b/api/core/tools/provider/builtin/webscraper/tools/webscraper.yaml @@ -0,0 +1,60 @@ +identity: + name: webscraper + author: Dify + label: + en_US: Web Scraper + zh_Hans: 网页爬虫 + pt_BR: Web Scraper +description: + human: + en_US: A tool for scraping webpages. + zh_Hans: 一个用于爬取网页的工具。 + pt_BR: A tool for scraping webpages. + llm: A tool for scraping webpages. Input should be a URL. +parameters: + - name: url + type: string + required: true + label: + en_US: URL + zh_Hans: 网页链接 + pt_BR: URL + human_description: + en_US: used for linking to webpages + zh_Hans: 用于链接到网页 + pt_BR: used for linking to webpages + llm_description: url for scraping + form: llm + - name: user_agent + type: string + required: false + label: + en_US: User Agent + zh_Hans: User Agent + pt_BR: User Agent + human_description: + en_US: used for identifying the browser. + zh_Hans: 用于识别浏览器。 + pt_BR: used for identifying the browser. + form: form + default: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/100.0.1000.0 Safari/537.36 + - name: generate_summary + type: boolean + required: false + label: + en_US: Whether to generate summary + zh_Hans: 是否生成摘要 + human_description: + en_US: If true, the crawler will only return the page summary content. + zh_Hans: 如果启用,爬虫将仅返回页面摘要内容。 + form: form + options: + - value: 'true' + label: + en_US: 'Yes' + zh_Hans: 是 + - value: 'false' + label: + en_US: 'No' + zh_Hans: 否 + default: 'false' diff --git a/api/core/tools/provider/builtin/webscraper/webscraper.py b/api/core/tools/provider/builtin/webscraper/webscraper.py new file mode 100644 index 0000000000000000000000000000000000000000..3c51393ac64cc4cfccce2c076598e098031d349c --- /dev/null +++ b/api/core/tools/provider/builtin/webscraper/webscraper.py @@ -0,0 +1,23 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.webscraper.tools.webscraper import WebscraperTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class WebscraperProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + WebscraperTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "url": "https://www.google.com", + "user_agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 ", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/webscraper/webscraper.yaml b/api/core/tools/provider/builtin/webscraper/webscraper.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6c2eb97784e2987744f6bf5edf3da92908eaf367 --- /dev/null +++ b/api/core/tools/provider/builtin/webscraper/webscraper.yaml @@ -0,0 +1,15 @@ +identity: + author: Dify + name: webscraper + label: + en_US: WebScraper + zh_Hans: 网页抓取 + pt_BR: WebScraper + description: + en_US: Web Scrapper tool kit is used to scrape web + zh_Hans: 一个用于抓取网页的工具。 + pt_BR: Web Scrapper tool kit is used to scrape web + icon: icon.svg + tags: + - productivity +credentials_for_provider: diff --git a/api/core/tools/provider/builtin/websearch/_assets/icon.svg b/api/core/tools/provider/builtin/websearch/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..d6ef5d878f863695179e852a3662564322fc5fa7 --- /dev/null +++ b/api/core/tools/provider/builtin/websearch/_assets/icon.svg @@ -0,0 +1,23 @@ + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/websearch/tools/get_markdown.py b/api/core/tools/provider/builtin/websearch/tools/get_markdown.py new file mode 100644 index 0000000000000000000000000000000000000000..043879deeab18f50ca15f984298d7cadcacc5504 --- /dev/null +++ b/api/core/tools/provider/builtin/websearch/tools/get_markdown.py @@ -0,0 +1,51 @@ +from typing import Any, Union + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +BASE_URL = "https://api.serply.io/v1/request" + + +class SerplyApi: + """ + SerplyAPI tool provider. + """ + + def __init__(self, api_key: str) -> None: + """Initialize SerplyAPI tool provider.""" + self.serply_api_key = api_key + + def run(self, url: str, **kwargs: Any) -> str: + """Run query through SerplyAPI and parse result.""" + + location = kwargs.get("location", "US") + + headers = { + "X-API-KEY": self.serply_api_key, + "X-User-Agent": kwargs.get("device", "desktop"), + "X-Proxy-Location": location, + "User-Agent": "Dify", + } + data = {"url": url, "method": "GET", "response_type": "markdown"} + res = requests.post(url, headers=headers, json=data) + return res.text + + +class GetMarkdownTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Invoke the SerplyApi tool. + """ + url = tool_parameters["url"] + location = tool_parameters.get("location") + + api_key = self.runtime.credentials["serply_api_key"] + result = SerplyApi(api_key).run(url, location=location) + + return self.create_text_message(text=result) diff --git a/api/core/tools/provider/builtin/websearch/tools/get_markdown.yaml b/api/core/tools/provider/builtin/websearch/tools/get_markdown.yaml new file mode 100644 index 0000000000000000000000000000000000000000..06a302bd14b82d11b22b9a7de17cfbfcf32cd4bb --- /dev/null +++ b/api/core/tools/provider/builtin/websearch/tools/get_markdown.yaml @@ -0,0 +1,96 @@ +identity: + name: get_markdown + author: Dify + label: + en_US: Get Markdown API + zh_Hans: Get Markdown API +description: + human: + en_US: A tool to perform convert a webpage to markdown to make it easier for LLMs to understand. + zh_Hans: 一个将网页转换为 Markdown 的工具,以便模型更容易理解 + llm: A tool to perform convert a webpage to markdown to make it easier for LLMs to understand. +parameters: + - name: url + type: string + required: true + label: + en_US: URL + zh_Hans: URL + human_description: + en_US: URL that you want to grab the content from + zh_Hans: 您要从中获取内容的 URL + llm_description: Defines the link want to grab content from. + form: llm + - name: location + type: string + required: false + default: US + label: + en_US: Location + zh_Hans: 询问 + human_description: + en_US: Defines from where you want the search to originate. (For example - New York) + zh_Hans: 定义您想要搜索的起始位置。 (例如 - 纽约) + llm_description: Defines from where you want the search to originate. (For example - New York) + form: form + options: + - value: AU + label: + en_US: Australia + zh_Hans: 澳大利亚 + pt_BR: Australia + - value: BR + label: + en_US: Brazil + zh_Hans: 巴西 + pt_BR: Brazil + - value: CA + label: + en_US: Canada + zh_Hans: 加拿大 + pt_BR: Canada + - value: DE + label: + en_US: Germany + zh_Hans: 德国 + pt_BR: Germany + - value: FR + label: + en_US: France + zh_Hans: 法国 + pt_BR: France + - value: GB + label: + en_US: United Kingdom + zh_Hans: 英国 + pt_BR: United Kingdom + - value: US + label: + en_US: United States + zh_Hans: 美国 + pt_BR: United States + - value: JP + label: + en_US: Japan + zh_Hans: 日本 + pt_BR: Japan + - value: IN + label: + en_US: India + zh_Hans: 印度 + pt_BR: India + - value: KR + label: + en_US: Korea + zh_Hans: 韩国 + pt_BR: Korea + - value: SG + label: + en_US: Singapore + zh_Hans: 新加坡 + pt_BR: Singapore + - value: SE + label: + en_US: Sweden + zh_Hans: 瑞典 + pt_BR: Sweden diff --git a/api/core/tools/provider/builtin/websearch/tools/job_search.py b/api/core/tools/provider/builtin/websearch/tools/job_search.py new file mode 100644 index 0000000000000000000000000000000000000000..13eb40339153c9995cbb13e62e24f061c773cfbd --- /dev/null +++ b/api/core/tools/provider/builtin/websearch/tools/job_search.py @@ -0,0 +1,88 @@ +from typing import Any, Union +from urllib.parse import urlencode + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +BASE_URL = "https://api.serply.io/v1/news/" + + +class SerplyApi: + """ + SerplyAPI tool provider. + """ + + def __init__(self, api_key: str) -> None: + """Initialize SerplyAPI tool provider.""" + self.serply_api_key = api_key + + def run(self, query: str, **kwargs: Any) -> str: + """Run query through SerplyAPI and parse result.""" + params = {"q": query, "hl": kwargs.get("hl", "en"), "gl": kwargs.get("gl", "US"), "num": kwargs.get("num", 10)} + location = kwargs.get("location", "US") + + headers = { + "X-API-KEY": self.serply_api_key, + "X-User-Agent": kwargs.get("device", "desktop"), + "X-Proxy-Location": location, + "User-Agent": "Dify", + } + + url = f"{BASE_URL}{urlencode(params)}" + res = requests.get( + url, + headers=headers, + ) + res = res.json() + + return self.parse_results(res) + + @staticmethod + def parse_results(res: dict) -> str: + """Process response from Serply Job Search.""" + jobs = res.get("jobs", []) + if not res or "jobs" not in res: + raise ValueError(f"Got error from Serply: {res}") + + string = [] + for job in jobs[:10]: + try: + string.append( + "\n".join( + [ + f"Position: {job['position']}", + f"Employer: {job['employer']}", + f"Location: {job['location']}", + f"Link: {job['link']}", + f"""Highest: {", ".join(list(job["highlights"]))}""", + "---", + ] + ) + ) + except KeyError: + continue + + content = "\n".join(string) + return f"\nJobs results:\n {content}\n" + + +class JobSearchTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Invoke the SerplyApi tool. + """ + query = tool_parameters["query"] + gl = tool_parameters.get("gl", "us") + hl = tool_parameters.get("hl", "en") + location = tool_parameters.get("location") + + api_key = self.runtime.credentials["serply_api_key"] + result = SerplyApi(api_key).run(query, gl=gl, hl=hl, location=location) + + return self.create_text_message(text=result) diff --git a/api/core/tools/provider/builtin/websearch/tools/job_search.yaml b/api/core/tools/provider/builtin/websearch/tools/job_search.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b5ede3df46ab01f71fa7ad1bb9dfd880645256e5 --- /dev/null +++ b/api/core/tools/provider/builtin/websearch/tools/job_search.yaml @@ -0,0 +1,41 @@ +identity: + name: job_search + author: Dify + label: + en_US: Job Search API + zh_Hans: Job Search API +description: + human: + en_US: A tool to retrieve job titles, company names and description from Google Jobs engine. + zh_Hans: 一个从 Google 招聘引擎检索职位名称、公司名称和描述的工具。 + llm: A tool to retrieve job titles, company names and description from Google Jobs engine. +parameters: + - name: query + type: string + required: true + label: + en_US: Query + zh_Hans: 询问 + human_description: + en_US: Defines the query you want to search. + zh_Hans: 定义您要搜索的查询。 + llm_description: Defines the search query you want to search. + form: llm + - name: location + type: string + required: false + default: US + label: + en_US: Location + zh_Hans: 询问 + human_description: + en_US: Defines from where you want the search to originate. (For example - New York) + zh_Hans: 定义您想要搜索的起始位置。 (例如 - 纽约) + llm_description: Defines from where you want the search to originate. (For example - New York) + form: form + options: + - value: US + label: + en_US: United States + zh_Hans: 美国 + pt_BR: United States diff --git a/api/core/tools/provider/builtin/websearch/tools/news_search.py b/api/core/tools/provider/builtin/websearch/tools/news_search.py new file mode 100644 index 0000000000000000000000000000000000000000..7a8a732ff3f2460e8eb9e777de3d8625401be244 --- /dev/null +++ b/api/core/tools/provider/builtin/websearch/tools/news_search.py @@ -0,0 +1,90 @@ +from typing import Any, Union +from urllib.parse import urlencode + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +BASE_URL = "https://api.serply.io/v1/news/" + + +class SerplyApi: + """ + SerplyApi tool provider. + """ + + def __init__(self, api_key: str) -> None: + """Initialize SerplyApi tool provider.""" + self.serply_api_key = api_key + + def run(self, query: str, **kwargs: Any) -> str: + """Run query through SerplyApi and parse result.""" + params = {"q": query, "hl": kwargs.get("hl", "en"), "gl": kwargs.get("gl", "US"), "num": kwargs.get("num", 10)} + location = kwargs.get("location", "US") + + headers = { + "X-API-KEY": self.serply_api_key, + "X-User-Agent": kwargs.get("device", "desktop"), + "X-Proxy-Location": location, + "User-Agent": "Dify", + } + + url = f"{BASE_URL}{urlencode(params)}" + res = requests.get( + url, + headers=headers, + ) + res = res.json() + + return self.parse_results(res) + + @staticmethod + def parse_results(res: dict) -> str: + """Process response from Serply News Search.""" + news = res.get("entries", []) + if not res or "entries" not in res: + raise ValueError(f"Got error from Serply: {res}") + + string = [] + for entry in news: + try: + # follow url + r = requests.get(entry["link"]) + final_link = r.history[-1].headers["Location"] + string.append( + "\n".join( + [ + f"Title: {entry['title']}", + f"Link: {final_link}", + f"Source: {entry['source']['title']}", + f"Published: {entry['published']}", + "---", + ] + ) + ) + except KeyError: + continue + + content = "\n".join(string) + return f"\nNews:\n {content}\n" + + +class NewsSearchTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Invoke the SerplyApi tool. + """ + query = tool_parameters["query"] + gl = tool_parameters.get("gl", "us") + hl = tool_parameters.get("hl", "en") + location = tool_parameters.get("location") + + api_key = self.runtime.credentials["serply_api_key"] + result = SerplyApi(api_key).run(query, gl=gl, hl=hl, location=location) + + return self.create_text_message(text=result) diff --git a/api/core/tools/provider/builtin/websearch/tools/news_search.yaml b/api/core/tools/provider/builtin/websearch/tools/news_search.yaml new file mode 100644 index 0000000000000000000000000000000000000000..126c610825ebbb63df91ea8a8547a39b3c81dff5 --- /dev/null +++ b/api/core/tools/provider/builtin/websearch/tools/news_search.yaml @@ -0,0 +1,501 @@ +identity: + name: news_search + author: Dify + label: + en_US: News Search API + zh_Hans: News Search API +description: + human: + en_US: A tool to retrieve organic search results snippets and links from Google News engine. + zh_Hans: 一种从 Google 新闻引擎检索有机搜索结果片段和链接的工具。 + llm: A tool to retrieve organic search results snippets and links from Google News engine. +parameters: + - name: query + type: string + required: true + label: + en_US: Query + zh_Hans: 询问 + human_description: + en_US: Defines the query you want to search. + zh_Hans: 定义您要搜索的查询。 + llm_description: Defines the search query you want to search. + form: llm + - name: location + type: string + required: false + default: US + label: + en_US: Location + zh_Hans: 询问 + human_description: + en_US: Defines from where you want the search to originate. (For example - New York) + zh_Hans: 定义您想要搜索的起始位置。 (例如 - 纽约) + llm_description: Defines from where you want the search to originate. (For example - New York) + form: form + options: + - value: AU + label: + en_US: Australia + zh_Hans: 澳大利亚 + pt_BR: Australia + - value: BR + label: + en_US: Brazil + zh_Hans: 巴西 + pt_BR: Brazil + - value: CA + label: + en_US: Canada + zh_Hans: 加拿大 + pt_BR: Canada + - value: DE + label: + en_US: Germany + zh_Hans: 德国 + pt_BR: Germany + - value: FR + label: + en_US: France + zh_Hans: 法国 + pt_BR: France + - value: GB + label: + en_US: United Kingdom + zh_Hans: 英国 + pt_BR: United Kingdom + - value: US + label: + en_US: United States + zh_Hans: 美国 + pt_BR: United States + - value: JP + label: + en_US: Japan + zh_Hans: 日本 + pt_BR: Japan + - value: IN + label: + en_US: India + zh_Hans: 印度 + pt_BR: India + - value: KR + label: + en_US: Korea + zh_Hans: 韩国 + pt_BR: Korea + - value: SG + label: + en_US: Singapore + zh_Hans: 新加坡 + pt_BR: Singapore + - value: SE + label: + en_US: Sweden + zh_Hans: 瑞典 + pt_BR: Sweden + - name: gl + type: select + label: + en_US: Country + zh_Hans: 国家/地区 + required: false + human_description: + en_US: Defines the country of the search. Default is "US". + zh_Hans: 定义搜索的国家/地区。默认为“美国”。 + llm_description: Defines the gl parameter of the Google search. + form: form + default: US + options: + - value: AR + label: + en_US: Argentina + zh_Hans: 阿根廷 + pt_BR: Argentina + - value: AU + label: + en_US: Australia + zh_Hans: 澳大利亚 + pt_BR: Australia + - value: AT + label: + en_US: Austria + zh_Hans: 奥地利 + pt_BR: Austria + - value: BE + label: + en_US: Belgium + zh_Hans: 比利时 + pt_BR: Belgium + - value: BR + label: + en_US: Brazil + zh_Hans: 巴西 + pt_BR: Brazil + - value: CA + label: + en_US: Canada + zh_Hans: 加拿大 + pt_BR: Canada + - value: CL + label: + en_US: Chile + zh_Hans: 智利 + pt_BR: Chile + - value: CO + label: + en_US: Colombia + zh_Hans: 哥伦比亚 + pt_BR: Colombia + - value: CN + label: + en_US: China + zh_Hans: 中国 + pt_BR: China + - value: CZ + label: + en_US: Czech Republic + zh_Hans: 捷克共和国 + pt_BR: Czech Republic + - value: DK + label: + en_US: Denmark + zh_Hans: 丹麦 + pt_BR: Denmark + - value: FI + label: + en_US: Finland + zh_Hans: 芬兰 + pt_BR: Finland + - value: FR + label: + en_US: France + zh_Hans: 法国 + pt_BR: France + - value: DE + label: + en_US: Germany + zh_Hans: 德国 + pt_BR: Germany + - value: HK + label: + en_US: Hong Kong + zh_Hans: 香港 + pt_BR: Hong Kong + - value: IN + label: + en_US: India + zh_Hans: 印度 + pt_BR: India + - value: ID + label: + en_US: Indonesia + zh_Hans: 印度尼西亚 + pt_BR: Indonesia + - value: IT + label: + en_US: Italy + zh_Hans: 意大利 + pt_BR: Italy + - value: JP + label: + en_US: Japan + zh_Hans: 日本 + pt_BR: Japan + - value: KR + label: + en_US: Korea + zh_Hans: 韩国 + pt_BR: Korea + - value: MY + label: + en_US: Malaysia + zh_Hans: 马来西亚 + pt_BR: Malaysia + - value: MX + label: + en_US: Mexico + zh_Hans: 墨西哥 + pt_BR: Mexico + - value: NL + label: + en_US: Netherlands + zh_Hans: 荷兰 + pt_BR: Netherlands + - value: NZ + label: + en_US: New Zealand + zh_Hans: 新西兰 + pt_BR: New Zealand + - value: NO + label: + en_US: Norway + zh_Hans: 挪威 + pt_BR: Norway + - value: PH + label: + en_US: Philippines + zh_Hans: 菲律宾 + pt_BR: Philippines + - value: PL + label: + en_US: Poland + zh_Hans: 波兰 + pt_BR: Poland + - value: PT + label: + en_US: Portugal + zh_Hans: 葡萄牙 + pt_BR: Portugal + - value: RU + label: + en_US: Russia + zh_Hans: 俄罗斯 + pt_BR: Russia + - value: SA + label: + en_US: Saudi Arabia + zh_Hans: 沙特阿拉伯 + pt_BR: Saudi Arabia + - value: SG + label: + en_US: Singapore + zh_Hans: 新加坡 + pt_BR: Singapore + - value: ZA + label: + en_US: South Africa + zh_Hans: 南非 + pt_BR: South Africa + - value: ES + label: + en_US: Spain + zh_Hans: 西班牙 + pt_BR: Spain + - value: SE + label: + en_US: Sweden + zh_Hans: 瑞典 + pt_BR: Sweden + - value: CH + label: + en_US: Switzerland + zh_Hans: 瑞士 + pt_BR: Switzerland + - value: TW + label: + en_US: Taiwan + zh_Hans: 台湾 + pt_BR: Taiwan + - value: TH + label: + en_US: Thailand + zh_Hans: 泰国 + pt_BR: Thailand + - value: TR + label: + en_US: Turkey + zh_Hans: 土耳其 + pt_BR: Turkey + - value: GB + label: + en_US: United Kingdom + zh_Hans: 英国 + pt_BR: United Kingdom + - value: US + label: + en_US: United States + zh_Hans: 美国 + pt_BR: United States + - name: hl + type: select + label: + en_US: Language + zh_Hans: 语言 + human_description: + en_US: Defines the interface language of the search. Default is "en". + zh_Hans: 定义搜索的界面语言。默认为“en”。 + required: false + default: en + form: form + options: + - value: ar + label: + en_US: Arabic + zh_Hans: 阿拉伯语 + - value: bg + label: + en_US: Bulgarian + zh_Hans: 保加利亚语 + - value: ca + label: + en_US: Catalan + zh_Hans: 加泰罗尼亚语 + - value: zh-cn + label: + en_US: Chinese (Simplified) + zh_Hans: 中文(简体) + - value: zh-tw + label: + en_US: Chinese (Traditional) + zh_Hans: 中文(繁体) + - value: cs + label: + en_US: Czech + zh_Hans: 捷克语 + - value: da + label: + en_US: Danish + zh_Hans: 丹麦语 + - value: nl + label: + en_US: Dutch + zh_Hans: 荷兰语 + - value: en + label: + en_US: English + zh_Hans: 英语 + - value: et + label: + en_US: Estonian + zh_Hans: 爱沙尼亚语 + - value: fi + label: + en_US: Finnish + zh_Hans: 芬兰语 + - value: fr + label: + en_US: French + zh_Hans: 法语 + - value: de + label: + en_US: German + zh_Hans: 德语 + - value: el + label: + en_US: Greek + zh_Hans: 希腊语 + - value: iw + label: + en_US: Hebrew + zh_Hans: 希伯来语 + - value: hi + label: + en_US: Hindi + zh_Hans: 印地语 + - value: hu + label: + en_US: Hungarian + zh_Hans: 匈牙利语 + - value: id + label: + en_US: Indonesian + zh_Hans: 印尼语 + - value: it + label: + en_US: Italian + zh_Hans: 意大利语 + - value: ja + label: + en_US: Japanese + zh_Hans: 日语 + - value: kn + label: + en_US: Kannada + zh_Hans: 卡纳达语 + - value: ko + label: + en_US: Korean + zh_Hans: 韩语 + - value: lv + label: + en_US: Latvian + zh_Hans: 拉脱维亚语 + - value: lt + label: + en_US: Lithuanian + zh_Hans: 立陶宛语 + - value: my + label: + en_US: Malay + zh_Hans: 马来语 + - value: ml + label: + en_US: Malayalam + zh_Hans: 马拉雅拉姆语 + - value: mr + label: + en_US: Marathi + zh_Hans: 马拉地语 + - value: "no" + label: + en_US: Norwegian + zh_Hans: 挪威语 + - value: pl + label: + en_US: Polish + zh_Hans: 波兰语 + - value: pt-br + label: + en_US: Portuguese (Brazil) + zh_Hans: 葡萄牙语(巴西) + - value: pt-pt + label: + en_US: Portuguese (Portugal) + zh_Hans: 葡萄牙语(葡萄牙) + - value: pa + label: + en_US: Punjabi + zh_Hans: 旁遮普语 + - value: ro + label: + en_US: Romanian + zh_Hans: 罗马尼亚语 + - value: ru + label: + en_US: Russian + zh_Hans: 俄语 + - value: sr + label: + en_US: Serbian + zh_Hans: 塞尔维亚语 + - value: sk + label: + en_US: Slovak + zh_Hans: 斯洛伐克语 + - value: sl + label: + en_US: Slovenian + zh_Hans: 斯洛文尼亚语 + - value: es + label: + en_US: Spanish + zh_Hans: 西班牙语 + - value: sv + label: + en_US: Swedish + zh_Hans: 瑞典语 + - value: ta + label: + en_US: Tamil + zh_Hans: 泰米尔语 + - value: te + label: + en_US: Telugu + zh_Hans: 泰卢固语 + - value: th + label: + en_US: Thai + zh_Hans: 泰语 + - value: tr + label: + en_US: Turkish + zh_Hans: 土耳其语 + - value: uk + label: + en_US: Ukrainian + zh_Hans: 乌克兰语 + - value: vi + label: + en_US: Vietnamese + zh_Hans: 越南语 diff --git a/api/core/tools/provider/builtin/websearch/tools/scholar_search.py b/api/core/tools/provider/builtin/websearch/tools/scholar_search.py new file mode 100644 index 0000000000000000000000000000000000000000..32c5d39e5b8674fada45afb7f063d75ca9252cb9 --- /dev/null +++ b/api/core/tools/provider/builtin/websearch/tools/scholar_search.py @@ -0,0 +1,93 @@ +from typing import Any, Union +from urllib.parse import urlencode + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +BASE_URL = "https://api.serply.io/v1/scholar/" + + +class SerplyApi: + """ + SerplyApi tool provider. + """ + + def __init__(self, api_key: str) -> None: + """Initialize SerplyApi tool provider.""" + self.serply_api_key = api_key + + def run(self, query: str, **kwargs: Any) -> str: + """Run query through SerplyApi and parse result.""" + params = {"q": query, "hl": kwargs.get("hl", "en"), "gl": kwargs.get("gl", "US"), "num": kwargs.get("num", 10)} + location = kwargs.get("location", "US") + + headers = { + "X-API-KEY": self.serply_api_key, + "X-User-Agent": kwargs.get("device", "desktop"), + "X-Proxy-Location": location, + "User-Agent": "Dify", + } + + url = f"{BASE_URL}{urlencode(params)}" + res = requests.get( + url, + headers=headers, + ) + res = res.json() + + return self.parse_results(res) + + @staticmethod + def parse_results(res: dict) -> str: + """Process response from Serply News Search.""" + articles = res.get("articles", []) + if not res or "articles" not in res: + raise ValueError(f"Got error from Serply: {res}") + + string = [] + for article in articles: + try: + if "doc" in article: + link = article["doc"]["link"] + else: + link = article["link"] + authors = [author["name"] for author in article["author"]["authors"]] + string.append( + "\n".join( + [ + f"Title: {article['title']}", + f"Link: {link}", + f"Description: {article['description']}", + f"Cite: {article['cite']}", + f"Authors: {', '.join(authors)}", + "---", + ] + ) + ) + except KeyError: + continue + + content = "\n".join(string) + return f"\nScholar results:\n {content}\n" + + +class ScholarSearchTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Invoke the SerplyApi tool. + """ + query = tool_parameters["query"] + gl = tool_parameters.get("gl", "us") + hl = tool_parameters.get("hl", "en") + location = tool_parameters.get("location") + + api_key = self.runtime.credentials["serply_api_key"] + result = SerplyApi(api_key).run(query, gl=gl, hl=hl, location=location) + + return self.create_text_message(text=result) diff --git a/api/core/tools/provider/builtin/websearch/tools/scholar_search.yaml b/api/core/tools/provider/builtin/websearch/tools/scholar_search.yaml new file mode 100644 index 0000000000000000000000000000000000000000..63e79d7ebfaa49448fe30255d3e129cd02fbfbea --- /dev/null +++ b/api/core/tools/provider/builtin/websearch/tools/scholar_search.yaml @@ -0,0 +1,501 @@ +identity: + name: scholar_search + author: Dify + label: + en_US: Scholar API + zh_Hans: Scholar API +description: + human: + en_US: A tool to retrieve scholarly literature. + zh_Hans: 学术文献检索工具 + llm: A tool to retrieve scholarly literature. +parameters: + - name: query + type: string + required: true + label: + en_US: Query + zh_Hans: 询问 + human_description: + en_US: Defines the query you want to search. + zh_Hans: 定义您要搜索的查询。 + llm_description: Defines the search query you want to search. + form: llm + - name: location + type: string + required: false + default: US + label: + en_US: Location + zh_Hans: 询问 + human_description: + en_US: Defines from where you want the search to originate. (For example - New York) + zh_Hans: 定义您想要搜索的起始位置。 (例如 - 纽约) + llm_description: Defines from where you want the search to originate. (For example - New York) + form: form + options: + - value: AU + label: + en_US: Australia + zh_Hans: 澳大利亚 + pt_BR: Australia + - value: BR + label: + en_US: Brazil + zh_Hans: 巴西 + pt_BR: Brazil + - value: CA + label: + en_US: Canada + zh_Hans: 加拿大 + pt_BR: Canada + - value: DE + label: + en_US: Germany + zh_Hans: 德国 + pt_BR: Germany + - value: FR + label: + en_US: France + zh_Hans: 法国 + pt_BR: France + - value: GB + label: + en_US: United Kingdom + zh_Hans: 英国 + pt_BR: United Kingdom + - value: US + label: + en_US: United States + zh_Hans: 美国 + pt_BR: United States + - value: JP + label: + en_US: Japan + zh_Hans: 日本 + pt_BR: Japan + - value: IN + label: + en_US: India + zh_Hans: 印度 + pt_BR: India + - value: KR + label: + en_US: Korea + zh_Hans: 韩国 + pt_BR: Korea + - value: SG + label: + en_US: Singapore + zh_Hans: 新加坡 + pt_BR: Singapore + - value: SE + label: + en_US: Sweden + zh_Hans: 瑞典 + pt_BR: Sweden + - name: gl + type: select + label: + en_US: Country + zh_Hans: 国家/地区 + required: false + human_description: + en_US: Defines the country of the search. Default is "US". + zh_Hans: 定义搜索的国家/地区。默认为“美国”。 + llm_description: Defines the gl parameter of the Google search. + form: form + default: US + options: + - value: AR + label: + en_US: Argentina + zh_Hans: 阿根廷 + pt_BR: Argentina + - value: AU + label: + en_US: Australia + zh_Hans: 澳大利亚 + pt_BR: Australia + - value: AT + label: + en_US: Austria + zh_Hans: 奥地利 + pt_BR: Austria + - value: BE + label: + en_US: Belgium + zh_Hans: 比利时 + pt_BR: Belgium + - value: BR + label: + en_US: Brazil + zh_Hans: 巴西 + pt_BR: Brazil + - value: CA + label: + en_US: Canada + zh_Hans: 加拿大 + pt_BR: Canada + - value: CL + label: + en_US: Chile + zh_Hans: 智利 + pt_BR: Chile + - value: CO + label: + en_US: Colombia + zh_Hans: 哥伦比亚 + pt_BR: Colombia + - value: CN + label: + en_US: China + zh_Hans: 中国 + pt_BR: China + - value: CZ + label: + en_US: Czech Republic + zh_Hans: 捷克共和国 + pt_BR: Czech Republic + - value: DK + label: + en_US: Denmark + zh_Hans: 丹麦 + pt_BR: Denmark + - value: FI + label: + en_US: Finland + zh_Hans: 芬兰 + pt_BR: Finland + - value: FR + label: + en_US: France + zh_Hans: 法国 + pt_BR: France + - value: DE + label: + en_US: Germany + zh_Hans: 德国 + pt_BR: Germany + - value: HK + label: + en_US: Hong Kong + zh_Hans: 香港 + pt_BR: Hong Kong + - value: IN + label: + en_US: India + zh_Hans: 印度 + pt_BR: India + - value: ID + label: + en_US: Indonesia + zh_Hans: 印度尼西亚 + pt_BR: Indonesia + - value: IT + label: + en_US: Italy + zh_Hans: 意大利 + pt_BR: Italy + - value: JP + label: + en_US: Japan + zh_Hans: 日本 + pt_BR: Japan + - value: KR + label: + en_US: Korea + zh_Hans: 韩国 + pt_BR: Korea + - value: MY + label: + en_US: Malaysia + zh_Hans: 马来西亚 + pt_BR: Malaysia + - value: MX + label: + en_US: Mexico + zh_Hans: 墨西哥 + pt_BR: Mexico + - value: NL + label: + en_US: Netherlands + zh_Hans: 荷兰 + pt_BR: Netherlands + - value: NZ + label: + en_US: New Zealand + zh_Hans: 新西兰 + pt_BR: New Zealand + - value: "NO" + label: + en_US: Norway + zh_Hans: 挪威 + pt_BR: Norway + - value: PH + label: + en_US: Philippines + zh_Hans: 菲律宾 + pt_BR: Philippines + - value: PL + label: + en_US: Poland + zh_Hans: 波兰 + pt_BR: Poland + - value: PT + label: + en_US: Portugal + zh_Hans: 葡萄牙 + pt_BR: Portugal + - value: RU + label: + en_US: Russia + zh_Hans: 俄罗斯 + pt_BR: Russia + - value: SA + label: + en_US: Saudi Arabia + zh_Hans: 沙特阿拉伯 + pt_BR: Saudi Arabia + - value: SG + label: + en_US: Singapore + zh_Hans: 新加坡 + pt_BR: Singapore + - value: ZA + label: + en_US: South Africa + zh_Hans: 南非 + pt_BR: South Africa + - value: ES + label: + en_US: Spain + zh_Hans: 西班牙 + pt_BR: Spain + - value: SE + label: + en_US: Sweden + zh_Hans: 瑞典 + pt_BR: Sweden + - value: CH + label: + en_US: Switzerland + zh_Hans: 瑞士 + pt_BR: Switzerland + - value: TW + label: + en_US: Taiwan + zh_Hans: 台湾 + pt_BR: Taiwan + - value: TH + label: + en_US: Thailand + zh_Hans: 泰国 + pt_BR: Thailand + - value: TR + label: + en_US: Turkey + zh_Hans: 土耳其 + pt_BR: Turkey + - value: GB + label: + en_US: United Kingdom + zh_Hans: 英国 + pt_BR: United Kingdom + - value: US + label: + en_US: United States + zh_Hans: 美国 + pt_BR: United States + - name: hl + type: select + label: + en_US: Language + zh_Hans: 语言 + human_description: + en_US: Defines the interface language of the search. Default is "en". + zh_Hans: 定义搜索的界面语言。默认为“en”。 + required: false + default: en + form: form + options: + - value: ar + label: + en_US: Arabic + zh_Hans: 阿拉伯语 + - value: bg + label: + en_US: Bulgarian + zh_Hans: 保加利亚语 + - value: ca + label: + en_US: Catalan + zh_Hans: 加泰罗尼亚语 + - value: zh-cn + label: + en_US: Chinese (Simplified) + zh_Hans: 中文(简体) + - value: zh-tw + label: + en_US: Chinese (Traditional) + zh_Hans: 中文(繁体) + - value: cs + label: + en_US: Czech + zh_Hans: 捷克语 + - value: da + label: + en_US: Danish + zh_Hans: 丹麦语 + - value: nl + label: + en_US: Dutch + zh_Hans: 荷兰语 + - value: en + label: + en_US: English + zh_Hans: 英语 + - value: et + label: + en_US: Estonian + zh_Hans: 爱沙尼亚语 + - value: fi + label: + en_US: Finnish + zh_Hans: 芬兰语 + - value: fr + label: + en_US: French + zh_Hans: 法语 + - value: de + label: + en_US: German + zh_Hans: 德语 + - value: el + label: + en_US: Greek + zh_Hans: 希腊语 + - value: iw + label: + en_US: Hebrew + zh_Hans: 希伯来语 + - value: hi + label: + en_US: Hindi + zh_Hans: 印地语 + - value: hu + label: + en_US: Hungarian + zh_Hans: 匈牙利语 + - value: id + label: + en_US: Indonesian + zh_Hans: 印尼语 + - value: it + label: + en_US: Italian + zh_Hans: 意大利语 + - value: ja + label: + en_US: Japanese + zh_Hans: 日语 + - value: kn + label: + en_US: Kannada + zh_Hans: 卡纳达语 + - value: ko + label: + en_US: Korean + zh_Hans: 韩语 + - value: lv + label: + en_US: Latvian + zh_Hans: 拉脱维亚语 + - value: lt + label: + en_US: Lithuanian + zh_Hans: 立陶宛语 + - value: my + label: + en_US: Malay + zh_Hans: 马来语 + - value: ml + label: + en_US: Malayalam + zh_Hans: 马拉雅拉姆语 + - value: mr + label: + en_US: Marathi + zh_Hans: 马拉地语 + - value: "no" + label: + en_US: Norwegian + zh_Hans: 挪威语 + - value: pl + label: + en_US: Polish + zh_Hans: 波兰语 + - value: pt-br + label: + en_US: Portuguese (Brazil) + zh_Hans: 葡萄牙语(巴西) + - value: pt-pt + label: + en_US: Portuguese (Portugal) + zh_Hans: 葡萄牙语(葡萄牙) + - value: pa + label: + en_US: Punjabi + zh_Hans: 旁遮普语 + - value: ro + label: + en_US: Romanian + zh_Hans: 罗马尼亚语 + - value: ru + label: + en_US: Russian + zh_Hans: 俄语 + - value: sr + label: + en_US: Serbian + zh_Hans: 塞尔维亚语 + - value: sk + label: + en_US: Slovak + zh_Hans: 斯洛伐克语 + - value: sl + label: + en_US: Slovenian + zh_Hans: 斯洛文尼亚语 + - value: es + label: + en_US: Spanish + zh_Hans: 西班牙语 + - value: sv + label: + en_US: Swedish + zh_Hans: 瑞典语 + - value: ta + label: + en_US: Tamil + zh_Hans: 泰米尔语 + - value: te + label: + en_US: Telugu + zh_Hans: 泰卢固语 + - value: th + label: + en_US: Thai + zh_Hans: 泰语 + - value: tr + label: + en_US: Turkish + zh_Hans: 土耳其语 + - value: uk + label: + en_US: Ukrainian + zh_Hans: 乌克兰语 + - value: vi + label: + en_US: Vietnamese + zh_Hans: 越南语 diff --git a/api/core/tools/provider/builtin/websearch/tools/web_search.py b/api/core/tools/provider/builtin/websearch/tools/web_search.py new file mode 100644 index 0000000000000000000000000000000000000000..d0e93cb0fa5681c021bbc8461875119038868a7c --- /dev/null +++ b/api/core/tools/provider/builtin/websearch/tools/web_search.py @@ -0,0 +1,90 @@ +import typing +from urllib.parse import urlencode + +import requests + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class SerplyApi: + """ + SerplyApi tool provider. + """ + + def __init__(self, api_key: str) -> None: + """Initialize Serply Web Search Tool provider.""" + self.serply_api_key = api_key + self.base_url = "https://api.serply.io/v1/search/" + + def run(self, query: str, **kwargs: typing.Any) -> str: + """Run query through Serply and parse result.""" + params = {"q": query, "hl": kwargs.get("hl", "en"), "gl": kwargs.get("gl", "US"), "num": kwargs.get("num", 10)} + location = kwargs.get("location", "US") + + headers = { + "X-API-KEY": self.serply_api_key, + "X-User-Agent": kwargs.get("device", "desktop"), + "X-Proxy-Location": location, + "User-Agent": "Dify", + } + + url = f"{self.base_url}{urlencode(params)}" + res = requests.get( + url, + headers=headers, + ) + res = res.json() + + return self.parse_results(res) + + @staticmethod + def parse_results(res: dict) -> str: + """Process response from Serply Web Search.""" + results = res.get("results", []) + if not res or "results" not in res: + raise ValueError(f"Got error from Serply: {res}") + + string = [] + for result in results: + try: + string.append( + "\n".join( + [ + f"Title: {result['title']}", + f"Link: {result['link']}", + f"Description: {result['description'].strip()}", + "---", + ] + ) + ) + except KeyError: + continue + + if related_questions := res.get("related_questions", []): + string.append("---") + string.append("Related Questions: ") + string.append("\n".join(related_questions)) + + content = "\n".join(string) + return f"\nSearch results:\n {content}\n" + + +class WebSearchTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, typing.Any], + ) -> typing.Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + Invoke the SerplyApi tool. + """ + query = tool_parameters["query"] + num = tool_parameters.get("num", 10) + gl = tool_parameters.get("gl", "us") + hl = tool_parameters.get("hl", "en") + location = tool_parameters.get("location", "None") + + api_key = self.runtime.credentials["serply_api_key"] + result = SerplyApi(api_key).run(query=query, num=num, gl=gl, hl=hl, location=location) + return self.create_text_message(text=result) diff --git a/api/core/tools/provider/builtin/websearch/tools/web_search.yaml b/api/core/tools/provider/builtin/websearch/tools/web_search.yaml new file mode 100644 index 0000000000000000000000000000000000000000..055029253c1753947f29b74f31c90a11b8c119d2 --- /dev/null +++ b/api/core/tools/provider/builtin/websearch/tools/web_search.yaml @@ -0,0 +1,376 @@ +identity: + name: web_search + author: Dify + label: + en_US: Web Search API + zh_Hans: Web Search API +description: + human: + en_US: A tool to retrieve answer boxes, knowledge graphs, snippets, and webpages from Google Search engine. + zh_Hans: 一种从 Google 搜索引擎检索答案框、知识图、片段和网页的工具。 + llm: A tool to retrieve answer boxes, knowledge graphs, snippets, and webpages from Google Search engine. +parameters: + - name: query + type: string + required: true + label: + en_US: Query + zh_Hans: 询问 + human_description: + en_US: Defines the query you want to search. + zh_Hans: 定义您要搜索的查询。 + llm_description: Defines the search query you want to search. + form: llm + - name: location + type: string + required: false + default: US + label: + en_US: Location + zh_Hans: 询问 + human_description: + en_US: Defines from where you want the search to originate. (For example - New York) + zh_Hans: 定义您想要搜索的起始位置。 (例如 - 纽约) + llm_description: Defines from where you want the search to originate. (For example - New York) + form: form + options: + - value: AU + label: + en_US: Australia + zh_Hans: 澳大利亚 + pt_BR: Australia + - value: BR + label: + en_US: Brazil + zh_Hans: 巴西 + pt_BR: Brazil + - value: CA + label: + en_US: Canada + zh_Hans: 加拿大 + pt_BR: Canada + - value: DE + label: + en_US: Germany + zh_Hans: 德国 + pt_BR: Germany + - value: FR + label: + en_US: France + zh_Hans: 法国 + pt_BR: France + - value: GB + label: + en_US: United Kingdom + zh_Hans: 英国 + pt_BR: United Kingdom + - value: US + label: + en_US: United States + zh_Hans: 美国 + pt_BR: United States + - value: JP + label: + en_US: Japan + zh_Hans: 日本 + pt_BR: Japan + - value: IN + label: + en_US: India + zh_Hans: 印度 + pt_BR: India + - value: KR + label: + en_US: Korea + zh_Hans: 韩国 + pt_BR: Korea + - value: SG + label: + en_US: Singapore + zh_Hans: 新加坡 + pt_BR: Singapore + - value: SE + label: + en_US: Sweden + zh_Hans: 瑞典 + pt_BR: Sweden + - name: device + type: select + label: + en_US: Device Type + zh_Hans: 汉斯先生 + human_description: + en_US: Defines the device to make interface search. Default is "desktop". + zh_Hans: 定义进行接口搜索的设备。默认为“桌面” + required: false + default: desktop + form: form + options: + - value: desktop + label: + en_US: Desktop + zh_Hans: 桌面 + - value: mobile + label: + en_US: Mobile + zh_Hans: 移动的 + - name: gl + type: select + label: + en_US: Country + zh_Hans: 国家/地区 + required: false + human_description: + en_US: Defines the country of the search. Default is "US". + zh_Hans: 定义搜索的国家/地区。默认为“美国”。 + llm_description: Defines the gl parameter of the Google search. + form: form + default: US + options: + - value: AU + label: + en_US: Australia + zh_Hans: 澳大利亚 + pt_BR: Australia + - value: BR + label: + en_US: Brazil + zh_Hans: 巴西 + pt_BR: Brazil + - value: CA + label: + en_US: Canada + zh_Hans: 加拿大 + pt_BR: Canada + - value: DE + label: + en_US: Germany + zh_Hans: 德国 + pt_BR: Germany + - value: FR + label: + en_US: France + zh_Hans: 法国 + pt_BR: France + - value: GB + label: + en_US: United Kingdom + zh_Hans: 英国 + pt_BR: United Kingdom + - value: IN + label: + en_US: India + zh_Hans: 印度 + pt_BR: India + - value: KR + label: + en_US: Korea + zh_Hans: 韩国 + pt_BR: Korea + - value: SE + label: + en_US: Sweden + zh_Hans: 瑞典 + pt_BR: Sweden + - value: SG + label: + en_US: Singapore + zh_Hans: 新加坡 + pt_BR: Singapore + - value: US + label: + en_US: United States + zh_Hans: 美国 + pt_BR: United States + - name: hl + type: select + label: + en_US: Language + zh_Hans: 语言 + human_description: + en_US: Defines the interface language of the search. Default is "en". + zh_Hans: 定义搜索的界面语言。默认为“en”。 + required: false + default: en + form: form + options: + - value: ar + label: + en_US: Arabic + zh_Hans: 阿拉伯语 + - value: bg + label: + en_US: Bulgarian + zh_Hans: 保加利亚语 + - value: ca + label: + en_US: Catalan + zh_Hans: 加泰罗尼亚语 + - value: zh-cn + label: + en_US: Chinese (Simplified) + zh_Hans: 中文(简体) + - value: zh-tw + label: + en_US: Chinese (Traditional) + zh_Hans: 中文(繁体) + - value: cs + label: + en_US: Czech + zh_Hans: 捷克语 + - value: da + label: + en_US: Danish + zh_Hans: 丹麦语 + - value: nl + label: + en_US: Dutch + zh_Hans: 荷兰语 + - value: en + label: + en_US: English + zh_Hans: 英语 + - value: et + label: + en_US: Estonian + zh_Hans: 爱沙尼亚语 + - value: fi + label: + en_US: Finnish + zh_Hans: 芬兰语 + - value: fr + label: + en_US: French + zh_Hans: 法语 + - value: de + label: + en_US: German + zh_Hans: 德语 + - value: el + label: + en_US: Greek + zh_Hans: 希腊语 + - value: iw + label: + en_US: Hebrew + zh_Hans: 希伯来语 + - value: hi + label: + en_US: Hindi + zh_Hans: 印地语 + - value: hu + label: + en_US: Hungarian + zh_Hans: 匈牙利语 + - value: id + label: + en_US: Indonesian + zh_Hans: 印尼语 + - value: it + label: + en_US: Italian + zh_Hans: 意大利语 + - value: ja + label: + en_US: Japanese + zh_Hans: 日语 + - value: kn + label: + en_US: Kannada + zh_Hans: 卡纳达语 + - value: ko + label: + en_US: Korean + zh_Hans: 韩语 + - value: lv + label: + en_US: Latvian + zh_Hans: 拉脱维亚语 + - value: lt + label: + en_US: Lithuanian + zh_Hans: 立陶宛语 + - value: my + label: + en_US: Malay + zh_Hans: 马来语 + - value: ml + label: + en_US: Malayalam + zh_Hans: 马拉雅拉姆语 + - value: mr + label: + en_US: Marathi + zh_Hans: 马拉地语 + - value: "no" + label: + en_US: Norwegian + zh_Hans: 挪威语 + - value: pl + label: + en_US: Polish + zh_Hans: 波兰语 + - value: pt-br + label: + en_US: Portuguese (Brazil) + zh_Hans: 葡萄牙语(巴西) + - value: pt-pt + label: + en_US: Portuguese (Portugal) + zh_Hans: 葡萄牙语(葡萄牙) + - value: pa + label: + en_US: Punjabi + zh_Hans: 旁遮普语 + - value: ro + label: + en_US: Romanian + zh_Hans: 罗马尼亚语 + - value: ru + label: + en_US: Russian + zh_Hans: 俄语 + - value: sr + label: + en_US: Serbian + zh_Hans: 塞尔维亚语 + - value: sk + label: + en_US: Slovak + zh_Hans: 斯洛伐克语 + - value: sl + label: + en_US: Slovenian + zh_Hans: 斯洛文尼亚语 + - value: es + label: + en_US: Spanish + zh_Hans: 西班牙语 + - value: sv + label: + en_US: Swedish + zh_Hans: 瑞典语 + - value: ta + label: + en_US: Tamil + zh_Hans: 泰米尔语 + - value: te + label: + en_US: Telugu + zh_Hans: 泰卢固语 + - value: th + label: + en_US: Thai + zh_Hans: 泰语 + - value: tr + label: + en_US: Turkish + zh_Hans: 土耳其语 + - value: uk + label: + en_US: Ukrainian + zh_Hans: 乌克兰语 + - value: vi + label: + en_US: Vietnamese + zh_Hans: 越南语 diff --git a/api/core/tools/provider/builtin/websearch/websearch.py b/api/core/tools/provider/builtin/websearch/websearch.py new file mode 100644 index 0000000000000000000000000000000000000000..90cc0c573ac97e23c82bec2c0c7cbb20420c4b96 --- /dev/null +++ b/api/core/tools/provider/builtin/websearch/websearch.py @@ -0,0 +1,21 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.websearch.tools.web_search import WebSearchTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class WebSearchAPIProvider(BuiltinToolProviderController): + # validate when saving the api_key + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + WebSearchTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={"query": "what is llm"}, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/websearch/websearch.yaml b/api/core/tools/provider/builtin/websearch/websearch.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c4267e1022dfa1572b6fddb1c7365947d74aa9eb --- /dev/null +++ b/api/core/tools/provider/builtin/websearch/websearch.yaml @@ -0,0 +1,34 @@ +identity: + name: websearch + author: Serply.io + label: + en_US: Serply.io + zh_Hans: Serply.io + pt_BR: Serply.io + description: + en_US: Serply.io is a robust real-time SERP API delivering structured data from a collection of search engines including Web Search, Jobs, News, and many more. + zh_Hans: Serply.io 是一个强大的实时 SERP API,可提供来自 搜索 招聘 新闻等搜索引擎集合的结构化数据。 + pt_BR: Serply.io is a robust real-time SERP API delivering structured data from a collection of search engines including Web Search, Jobs, News, and many more. + icon: icon.svg + tags: + - search + - business + - news + - productivity +credentials_for_provider: + serply_api_key: + type: secret-input + required: true + label: + en_US: Serply.io API key + zh_Hans: Serply.io API key + pt_BR: Serply.io API key + placeholder: + en_US: Please input your Serply.io API key + zh_Hans: 请输入你的 Serply.io API key + pt_BR: Please input your Serply.io API key + help: + en_US: Get your Serply.io API key from https://Serply.io/ + zh_Hans: 从 Serply.io 获取您的 Serply.io API key + pt_BR: Get your Serply.io API key from Serply.io + url: https://Serply.io/ diff --git a/api/core/tools/provider/builtin/wecom/_assets/icon.png b/api/core/tools/provider/builtin/wecom/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..6cc752af480217d97eee7071bc8e23cc9425ebc2 --- /dev/null +++ b/api/core/tools/provider/builtin/wecom/_assets/icon.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1cc1cd3129a939539d6a5a511066061b7f6608c319016ba77f86e4378c352691 +size 262939 diff --git a/api/core/tools/provider/builtin/wecom/tools/wecom_group_bot.py b/api/core/tools/provider/builtin/wecom/tools/wecom_group_bot.py new file mode 100644 index 0000000000000000000000000000000000000000..545d9f4f8d6335497e7316b0740a0da8e512a6c0 --- /dev/null +++ b/api/core/tools/provider/builtin/wecom/tools/wecom_group_bot.py @@ -0,0 +1,57 @@ +from typing import Any, Union + +import httpx + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.utils.uuid_utils import is_valid_uuid + + +class WecomGroupBotTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + content = tool_parameters.get("content", "") + if not content: + return self.create_text_message("Invalid parameter content") + + hook_key = tool_parameters.get("hook_key", "") + if not is_valid_uuid(hook_key): + return self.create_text_message(f"Invalid parameter hook_key ${hook_key}, not a valid UUID") + + message_type = tool_parameters.get("message_type", "text") + if message_type == "markdown": + payload = { + "msgtype": "markdown", + "markdown": { + "content": content, + }, + } + else: + payload = { + "msgtype": "text", + "text": { + "content": content, + }, + } + api_url = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send" + headers = { + "Content-Type": "application/json", + } + params = { + "key": hook_key, + } + + try: + res = httpx.post(api_url, headers=headers, params=params, json=payload) + if res.is_success: + return self.create_text_message("Text message sent successfully") + else: + return self.create_text_message( + f"Failed to send the text message, status code: {res.status_code}, response: {res.text}" + ) + except Exception as e: + return self.create_text_message("Failed to send message to group chat bot. {}".format(e)) diff --git a/api/core/tools/provider/builtin/wecom/tools/wecom_group_bot.yaml b/api/core/tools/provider/builtin/wecom/tools/wecom_group_bot.yaml new file mode 100644 index 0000000000000000000000000000000000000000..379005a10214200f3c62987a51c6ae206b6a7812 --- /dev/null +++ b/api/core/tools/provider/builtin/wecom/tools/wecom_group_bot.yaml @@ -0,0 +1,64 @@ +identity: + name: wecom_group_bot + author: Bowen Liang + label: + en_US: Send Group Message + zh_Hans: 发送群消息 + pt_BR: Send Group Message + icon: icon.svg +description: + human: + en_US: Sending a group message on Wecom via the webhook of group bot + zh_Hans: 通过企业微信的群机器人webhook发送群消息 + pt_BR: Sending a group message on Wecom via the webhook of group bot + llm: A tool for sending messages to a chat group on Wecom(企业微信) . +parameters: + - name: hook_key + type: secret-input + required: true + label: + en_US: Wecom Group bot webhook key + zh_Hans: 群机器人webhook的key + pt_BR: Wecom Group bot webhook key + human_description: + en_US: Wecom Group bot webhook key + zh_Hans: 群机器人webhook的key + pt_BR: Wecom Group bot webhook key + form: form + - name: content + type: string + required: true + label: + en_US: content + zh_Hans: 消息内容 + pt_BR: content + human_description: + en_US: Content to sent to the group. + zh_Hans: 群消息文本 + pt_BR: Content to sent to the group. + llm_description: Content of the message + form: llm + - name: message_type + type: select + default: text + required: true + label: + en_US: Wecom Group bot message type + zh_Hans: 群机器人webhook的消息类型 + pt_BR: Wecom Group bot message type + human_description: + en_US: Wecom Group bot message type + zh_Hans: 群机器人webhook的消息类型 + pt_BR: Wecom Group bot message type + options: + - value: text + label: + en_US: Text + zh_Hans: 文本 + pt_BR: Text + - value: markdown + label: + en_US: Markdown + zh_Hans: Markdown + pt_BR: Markdown + form: form diff --git a/api/core/tools/provider/builtin/wecom/wecom.py b/api/core/tools/provider/builtin/wecom/wecom.py new file mode 100644 index 0000000000000000000000000000000000000000..573f76ee56da67719b0891c41ce9c37d4df1eebe --- /dev/null +++ b/api/core/tools/provider/builtin/wecom/wecom.py @@ -0,0 +1,7 @@ +from core.tools.provider.builtin.wecom.tools.wecom_group_bot import WecomGroupBotTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class WecomProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + WecomGroupBotTool() diff --git a/api/core/tools/provider/builtin/wecom/wecom.yaml b/api/core/tools/provider/builtin/wecom/wecom.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a544055ba4cb6737ce5d278d40b0a7fd192c5088 --- /dev/null +++ b/api/core/tools/provider/builtin/wecom/wecom.yaml @@ -0,0 +1,15 @@ +identity: + author: Bowen Liang + name: wecom + label: + en_US: Wecom + zh_Hans: 企业微信 + pt_BR: Wecom + description: + en_US: Wecom group bot + zh_Hans: 企业微信群机器人 + pt_BR: Wecom group bot + icon: icon.png + tags: + - social +credentials_for_provider: diff --git a/api/core/tools/provider/builtin/wikipedia/_assets/icon.svg b/api/core/tools/provider/builtin/wikipedia/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..fe652aacf9c871f0832c910c567255540d62624b --- /dev/null +++ b/api/core/tools/provider/builtin/wikipedia/_assets/icon.svg @@ -0,0 +1,3 @@ + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/wikipedia/tools/wikipedia_search.py b/api/core/tools/provider/builtin/wikipedia/tools/wikipedia_search.py new file mode 100644 index 0000000000000000000000000000000000000000..edb96e722f7f335c0814ec69cdd7de9ffee36ac2 --- /dev/null +++ b/api/core/tools/provider/builtin/wikipedia/tools/wikipedia_search.py @@ -0,0 +1,105 @@ +from typing import Any, Optional, Union + +import wikipedia # type: ignore + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + +WIKIPEDIA_MAX_QUERY_LENGTH = 300 + + +class WikipediaAPIWrapper: + """Wrapper around WikipediaAPI. + + To use, you should have the ``wikipedia`` python package installed. + This wrapper will use the Wikipedia API to conduct searches and + fetch page summaries. By default, it will return the page summaries + of the top-k results. + It limits the Document content by doc_content_chars_max. + """ + + top_k_results: int = 3 + lang: str = "en" + load_all_available_meta: bool = False + doc_content_chars_max: int = 4000 + + def __init__(self, doc_content_chars_max: int = 4000): + self.doc_content_chars_max = doc_content_chars_max + + def run(self, query: str, lang: str = "") -> str: + if lang in wikipedia.languages(): + self.lang = lang + + wikipedia.set_lang(self.lang) + wiki_client = wikipedia + + """Run Wikipedia search and get page summaries.""" + page_titles = wiki_client.search(query[:WIKIPEDIA_MAX_QUERY_LENGTH]) + summaries = [] + for page_title in page_titles[: self.top_k_results]: + if wiki_page := self._fetch_page(page_title): + if summary := self._formatted_page_summary(page_title, wiki_page): + summaries.append(summary) + if not summaries: + return "No good Wikipedia Search Result was found" + return "\n\n".join(summaries)[: self.doc_content_chars_max] + + @staticmethod + def _formatted_page_summary(page_title: str, wiki_page: Any) -> Optional[str]: + return f"Page: {page_title}\nSummary: {wiki_page.summary}" + + def _fetch_page(self, page: str) -> Optional[str]: + try: + return wikipedia.page(title=page, auto_suggest=False) + except ( + wikipedia.exceptions.PageError, + wikipedia.exceptions.DisambiguationError, + ): + return None + + +class WikipediaQueryRun: + """Tool that searches the Wikipedia API.""" + + name = "Wikipedia" + description = ( + "A wrapper around Wikipedia. " + "Useful for when you need to answer general questions about " + "people, places, companies, facts, historical events, or other subjects. " + "Input should be a search query." + ) + api_wrapper: WikipediaAPIWrapper + + def __init__(self, api_wrapper: WikipediaAPIWrapper): + self.api_wrapper = api_wrapper + + def _run( + self, + query: str, + lang: str = "", + ) -> str: + """Use the Wikipedia tool.""" + return self.api_wrapper.run(query, lang) + + +class WikiPediaSearchTool(BuiltinTool): + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + query = tool_parameters.get("query", "") + lang = tool_parameters.get("language", "") + if not query: + return self.create_text_message("Please input query") + + tool = WikipediaQueryRun( + api_wrapper=WikipediaAPIWrapper(doc_content_chars_max=4000), + ) + + result = tool._run(query, lang) + + return self.create_text_message(self.summary(user_id=user_id, content=result)) diff --git a/api/core/tools/provider/builtin/wikipedia/tools/wikipedia_search.yaml b/api/core/tools/provider/builtin/wikipedia/tools/wikipedia_search.yaml new file mode 100644 index 0000000000000000000000000000000000000000..98d002df1c0daa666af16067a60cb79e63688b46 --- /dev/null +++ b/api/core/tools/provider/builtin/wikipedia/tools/wikipedia_search.yaml @@ -0,0 +1,101 @@ +identity: + name: wikipedia_search + author: Dify + label: + en_US: WikipediaSearch + zh_Hans: 维基百科搜索 + pt_BR: WikipediaSearch + icon: icon.svg +description: + human: + en_US: A tool for performing a Wikipedia search and extracting snippets and webpages. + zh_Hans: 一个用于执行维基百科搜索并提取片段和网页的工具。 + pt_BR: A tool for performing a Wikipedia search and extracting snippets and webpages. + llm: A tool for performing a Wikipedia search and extracting snippets and webpages. Input should be a search query. +parameters: + - name: query + type: string + required: true + label: + en_US: Query string + zh_Hans: 查询语句 + pt_BR: Query string + human_description: + en_US: key words for searching + zh_Hans: 查询关键词 + pt_BR: key words for searching + llm_description: key words for searching, this should be in the language of "language" parameter + form: llm + - name: language + type: string + required: true + label: + en_US: Language + zh_Hans: 语言 + human_description: + en_US: The language of the Wikipedia to be searched + zh_Hans: 要搜索的维基百科语言 + llm_description: >- + language of the wikipedia to be searched, + only "de" for German, + "en" for English, + "fr" for French, + "hi" for Hindi, + "ja" for Japanese, + "ko" for Korean, + "pl" for Polish, + "pt" for Portuguese, + "ro" for Romanian, + "uk" for Ukrainian, + "vi" for Vietnamese, + and "zh" for Chinese are supported + form: llm + options: + - value: de + label: + en_US: German + zh_Hans: 德语 + - value: en + label: + en_US: English + zh_Hans: 英语 + - value: fr + label: + en_US: French + zh_Hans: 法语 + - value: hi + label: + en_US: Hindi + zh_Hans: 印地语 + - value: ja + label: + en_US: Japanese + zh_Hans: 日语 + - value: ko + label: + en_US: Korean + zh_Hans: 韩语 + - value: pl + label: + en_US: Polish + zh_Hans: 波兰语 + - value: pt + label: + en_US: Portuguese + zh_Hans: 葡萄牙语 + - value: ro + label: + en_US: Romanian + zh_Hans: 罗马尼亚语 + - value: uk + label: + en_US: Ukrainian + zh_Hans: 乌克兰语 + - value: vi + label: + en_US: Vietnamese + zh_Hans: 越南语 + - value: zh + label: + en_US: Chinese + zh_Hans: 中文 diff --git a/api/core/tools/provider/builtin/wikipedia/wikipedia.py b/api/core/tools/provider/builtin/wikipedia/wikipedia.py new file mode 100644 index 0000000000000000000000000000000000000000..178bf7b0ceb2e970c05276fd061dae4c34abdd4a --- /dev/null +++ b/api/core/tools/provider/builtin/wikipedia/wikipedia.py @@ -0,0 +1,20 @@ +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.wikipedia.tools.wikipedia_search import WikiPediaSearchTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class WikiPediaProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + try: + WikiPediaSearchTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "query": "misaka mikoto", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/wikipedia/wikipedia.yaml b/api/core/tools/provider/builtin/wikipedia/wikipedia.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c5828240225d007e8dc5e78564efc54570b36245 --- /dev/null +++ b/api/core/tools/provider/builtin/wikipedia/wikipedia.yaml @@ -0,0 +1,15 @@ +identity: + author: Dify + name: wikipedia + label: + en_US: Wikipedia + zh_Hans: 维基百科 + pt_BR: Wikipedia + description: + en_US: Wikipedia is a free online encyclopedia, created and edited by volunteers around the world. + zh_Hans: 维基百科是一个由全世界的志愿者创建和编辑的免费在线百科全书。 + pt_BR: Wikipedia is a free online encyclopedia, created and edited by volunteers around the world. + icon: icon.svg + tags: + - social +credentials_for_provider: diff --git a/api/core/tools/provider/builtin/wolframalpha/_assets/icon.svg b/api/core/tools/provider/builtin/wolframalpha/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..2caf32ee67be0a8240d4d03fc886b7c68e20992a --- /dev/null +++ b/api/core/tools/provider/builtin/wolframalpha/_assets/icon.svg @@ -0,0 +1,23 @@ + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/wolframalpha/tools/wolframalpha.py b/api/core/tools/provider/builtin/wolframalpha/tools/wolframalpha.py new file mode 100644 index 0000000000000000000000000000000000000000..9b24be7cab81eb998c2342ecd9071f150b70703a --- /dev/null +++ b/api/core/tools/provider/builtin/wolframalpha/tools/wolframalpha.py @@ -0,0 +1,72 @@ +from typing import Any, Union + +from httpx import get + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.errors import ToolInvokeError, ToolProviderCredentialValidationError +from core.tools.tool.builtin_tool import BuiltinTool + + +class WolframAlphaTool(BuiltinTool): + _base_url = "https://api.wolframalpha.com/v2/query" + + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + query = tool_parameters.get("query", "") + if not query: + return self.create_text_message("Please input query") + appid = self.runtime.credentials.get("appid", "") + if not appid: + raise ToolProviderCredentialValidationError("Please input appid") + + params = {"appid": appid, "input": query, "includepodid": "Result", "format": "plaintext", "output": "json"} + + finished = False + result = None + # try 3 times at most + counter = 0 + + while not finished and counter < 3: + counter += 1 + try: + response = get(self._base_url, params=params, timeout=20) + response.raise_for_status() + response_data = response.json() + except Exception as e: + raise ToolInvokeError(str(e)) + + if "success" not in response_data["queryresult"] or response_data["queryresult"]["success"] != True: + query_result = response_data.get("queryresult", {}) + if query_result.get("error"): + if "msg" in query_result["error"]: + if query_result["error"]["msg"] == "Invalid appid": + raise ToolProviderCredentialValidationError("Invalid appid") + raise ToolInvokeError("Failed to invoke tool") + + if "didyoumeans" in response_data["queryresult"]: + # get the most likely interpretation + query = "" + max_score = 0 + for didyoumean in response_data["queryresult"]["didyoumeans"]: + if float(didyoumean["score"]) > max_score: + query = didyoumean["val"] + max_score = float(didyoumean["score"]) + + params["input"] = query + else: + finished = True + if "sources" in response_data["queryresult"]: + return self.create_link_message(response_data["queryresult"]["sources"]["url"]) + elif "pods" in response_data["queryresult"]: + result = response_data["queryresult"]["pods"][0]["subpods"][0]["plaintext"] + + if not finished or not result: + return self.create_text_message("No result found") + + return self.create_text_message(result) diff --git a/api/core/tools/provider/builtin/wolframalpha/tools/wolframalpha.yaml b/api/core/tools/provider/builtin/wolframalpha/tools/wolframalpha.yaml new file mode 100644 index 0000000000000000000000000000000000000000..08b5668691e23ac6e83948eb7d2dee6c299b3e71 --- /dev/null +++ b/api/core/tools/provider/builtin/wolframalpha/tools/wolframalpha.yaml @@ -0,0 +1,27 @@ +identity: + name: wolframalpha + author: Dify + label: + en_US: WolframAlpha + zh_Hans: WolframAlpha + pt_BR: WolframAlpha +description: + human: + en_US: WolframAlpha is a powerful computational knowledge engine. + zh_Hans: WolframAlpha 是一个强大的计算知识引擎。 + pt_BR: WolframAlpha is a powerful computational knowledge engine. + llm: WolframAlpha is a powerful computational knowledge engine. one single query can get the answer of a question. +parameters: + - name: query + type: string + required: true + label: + en_US: Query string + zh_Hans: 计算语句 + pt_BR: Query string + human_description: + en_US: used for calculating + zh_Hans: 用于计算最终结果 + pt_BR: used for calculating + llm_description: a single query for calculating + form: llm diff --git a/api/core/tools/provider/builtin/wolframalpha/wolframalpha.py b/api/core/tools/provider/builtin/wolframalpha/wolframalpha.py new file mode 100644 index 0000000000000000000000000000000000000000..7be288b5387f346f88a65c378788ba904da18ba7 --- /dev/null +++ b/api/core/tools/provider/builtin/wolframalpha/wolframalpha.py @@ -0,0 +1,22 @@ +from typing import Any + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.wolframalpha.tools.wolframalpha import WolframAlphaTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class GoogleProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + try: + WolframAlphaTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "query": "1+2+....+111", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/wolframalpha/wolframalpha.yaml b/api/core/tools/provider/builtin/wolframalpha/wolframalpha.yaml new file mode 100644 index 0000000000000000000000000000000000000000..91265eb3c00d0aaaca1c776e1531d0afea4bad5b --- /dev/null +++ b/api/core/tools/provider/builtin/wolframalpha/wolframalpha.yaml @@ -0,0 +1,32 @@ +identity: + author: Dify + name: wolframalpha + label: + en_US: WolframAlpha + zh_Hans: WolframAlpha + pt_BR: WolframAlpha + description: + en_US: WolframAlpha is a powerful computational knowledge engine. + zh_Hans: WolframAlpha 是一个强大的计算知识引擎。 + pt_BR: WolframAlpha is a powerful computational knowledge engine. + icon: icon.svg + tags: + - productivity + - utilities +credentials_for_provider: + appid: + type: secret-input + required: true + label: + en_US: WolframAlpha AppID + zh_Hans: WolframAlpha AppID + pt_BR: WolframAlpha AppID + placeholder: + en_US: Please input your WolframAlpha AppID + zh_Hans: 请输入你的 WolframAlpha AppID + pt_BR: Please input your WolframAlpha AppID + help: + en_US: Get your WolframAlpha AppID from WolframAlpha, please use "full results" api access. + zh_Hans: 从 WolframAlpha 获取您的 WolframAlpha AppID,请使用 "full results" API。 + pt_BR: Get your WolframAlpha AppID from WolframAlpha, please use "full results" api access. + url: https://products.wolframalpha.com/api diff --git a/api/core/tools/provider/builtin/xinference/_assets/icon.png b/api/core/tools/provider/builtin/xinference/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..e58cacbd123b5887b34fc8414d8b57aa801bb690 Binary files /dev/null and b/api/core/tools/provider/builtin/xinference/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/xinference/tools/stable_diffusion.py b/api/core/tools/provider/builtin/xinference/tools/stable_diffusion.py new file mode 100644 index 0000000000000000000000000000000000000000..a44d3b730a84f98f96bb62505d4395bd32ff395d --- /dev/null +++ b/api/core/tools/provider/builtin/xinference/tools/stable_diffusion.py @@ -0,0 +1,415 @@ +import io +import json +from base64 import b64decode, b64encode +from copy import deepcopy +from typing import Any, Union + +from httpx import get, post +from PIL import Image +from yarl import URL + +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_entities import ( + ToolInvokeMessage, + ToolParameter, + ToolParameterOption, +) +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.tool.builtin_tool import BuiltinTool + +# All commented out parameters default to null +DRAW_TEXT_OPTIONS = { + # Prompts + "prompt": "", + "negative_prompt": "", + # "styles": [], + # Seeds + "seed": -1, + "subseed": -1, + "subseed_strength": 0, + "seed_resize_from_h": -1, + "seed_resize_from_w": -1, + # Samplers + "sampler_name": "DPM++ 2M", + # "scheduler": "", + # "sampler_index": "Automatic", + # Latent Space Options + "batch_size": 1, + "n_iter": 1, + "steps": 10, + "cfg_scale": 7, + "width": 512, + "height": 512, + # "restore_faces": True, + # "tiling": True, + "do_not_save_samples": False, + "do_not_save_grid": False, + # "eta": 0, + # "denoising_strength": 0.75, + # "s_min_uncond": 0, + # "s_churn": 0, + # "s_tmax": 0, + # "s_tmin": 0, + # "s_noise": 0, + "override_settings": {}, + "override_settings_restore_afterwards": True, + # Refinement Options + "refiner_checkpoint": "", + "refiner_switch_at": 0, + "disable_extra_networks": False, + # "firstpass_image": "", + # "comments": "", + # High-Resolution Options + "enable_hr": False, + "firstphase_width": 0, + "firstphase_height": 0, + "hr_scale": 2, + # "hr_upscaler": "", + "hr_second_pass_steps": 0, + "hr_resize_x": 0, + "hr_resize_y": 0, + # "hr_checkpoint_name": "", + # "hr_sampler_name": "", + # "hr_scheduler": "", + "hr_prompt": "", + "hr_negative_prompt": "", + # Task Options + # "force_task_id": "", + # Script Options + # "script_name": "", + "script_args": [], + # Output Options + "send_images": True, + "save_images": False, + "alwayson_scripts": {}, + # "infotext": "", +} + + +class StableDiffusionTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + # base url + base_url = self.runtime.credentials.get("base_url", None) + if not base_url: + return self.create_text_message("Please input base_url") + + if tool_parameters.get("model"): + self.runtime.credentials["model"] = tool_parameters["model"] + + model = self.runtime.credentials.get("model", None) + if not model: + return self.create_text_message("Please input model") + api_key = self.runtime.credentials.get("api_key") or "abc" + headers = {"Authorization": f"Bearer {api_key}"} + # set model + try: + url = str(URL(base_url) / "sdapi" / "v1" / "options") + response = post( + url, + json={"sd_model_checkpoint": model}, + headers=headers, + ) + if response.status_code != 200: + raise ToolProviderCredentialValidationError("Failed to set model, please tell user to set model") + except Exception as e: + raise ToolProviderCredentialValidationError("Failed to set model, please tell user to set model") + + # get image id and image variable + image_id = tool_parameters.get("image_id", "") + image_variable = self.get_default_image_variable() + # Return text2img if there's no image ID or no image variable + if not image_id or not image_variable: + return self.text2img(base_url=base_url, tool_parameters=tool_parameters) + + # Proceed with image-to-image generation + return self.img2img(base_url=base_url, tool_parameters=tool_parameters) + + def validate_models(self): + """ + validate models + """ + try: + base_url = self.runtime.credentials.get("base_url", None) + if not base_url: + raise ToolProviderCredentialValidationError("Please input base_url") + model = self.runtime.credentials.get("model", None) + if not model: + raise ToolProviderCredentialValidationError("Please input model") + + api_url = str(URL(base_url) / "sdapi" / "v1" / "sd-models") + response = get(url=api_url, timeout=10) + if response.status_code == 404: + # try draw a picture + self._invoke( + user_id="test", + tool_parameters={ + "prompt": "a cat", + "width": 1024, + "height": 1024, + "steps": 1, + "lora": "", + }, + ) + elif response.status_code != 200: + raise ToolProviderCredentialValidationError("Failed to get models") + else: + models = [d["model_name"] for d in response.json()] + if len([d for d in models if d == model]) > 0: + return self.create_text_message(json.dumps(models)) + else: + raise ToolProviderCredentialValidationError(f"model {model} does not exist") + except Exception as e: + raise ToolProviderCredentialValidationError(f"Failed to get models, {e}") + + def get_sd_models(self) -> list[str]: + """ + get sd models + """ + try: + base_url = self.runtime.credentials.get("base_url", None) + if not base_url: + return [] + api_url = str(URL(base_url) / "sdapi" / "v1" / "sd-models") + response = get(url=api_url, timeout=120) + if response.status_code != 200: + return [] + else: + return [d["model_name"] for d in response.json()] + except Exception as e: + return [] + + def get_sample_methods(self) -> list[str]: + """ + get sample method + """ + try: + base_url = self.runtime.credentials.get("base_url", None) + if not base_url: + return [] + api_url = str(URL(base_url) / "sdapi" / "v1" / "samplers") + response = get(url=api_url, timeout=120) + if response.status_code != 200: + return [] + else: + return [d["name"] for d in response.json()] + except Exception as e: + return [] + + def img2img( + self, base_url: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + generate image + """ + + # Fetch the binary data of the image + image_variable = self.get_default_image_variable() + image_binary = self.get_variable_file(image_variable.name) + if not image_binary: + return self.create_text_message("Image not found, please request user to generate image firstly.") + + # Convert image to RGB and save as PNG + try: + with Image.open(io.BytesIO(image_binary)) as image, io.BytesIO() as buffer: + image.convert("RGB").save(buffer, format="PNG") + image_binary = buffer.getvalue() + except Exception as e: + return self.create_text_message(f"Failed to process the image: {str(e)}") + + # copy draw options + draw_options = deepcopy(DRAW_TEXT_OPTIONS) + # set image options + model = tool_parameters.get("model", "") + draw_options_image = { + "init_images": [b64encode(image_binary).decode("utf-8")], + "denoising_strength": 0.9, + "restore_faces": False, + "script_args": [], + "override_settings": {"sd_model_checkpoint": model}, + "resize_mode": 0, + "image_cfg_scale": 0, + # "mask": None, + "mask_blur_x": 4, + "mask_blur_y": 4, + "mask_blur": 0, + "mask_round": True, + "inpainting_fill": 0, + "inpaint_full_res": True, + "inpaint_full_res_padding": 0, + "inpainting_mask_invert": 0, + "initial_noise_multiplier": 0, + # "latent_mask": None, + "include_init_images": True, + } + # update key and values + draw_options.update(draw_options_image) + draw_options.update(tool_parameters) + + # get prompt lora model + prompt = tool_parameters.get("prompt", "") + lora = tool_parameters.get("lora", "") + model = tool_parameters.get("model", "") + if lora: + draw_options["prompt"] = f"{lora},{prompt}" + else: + draw_options["prompt"] = prompt + api_key = self.runtime.credentials.get("api_key") or "abc" + headers = {"Authorization": f"Bearer {api_key}"} + try: + url = str(URL(base_url) / "sdapi" / "v1" / "img2img") + response = post( + url, + json=draw_options, + timeout=120, + headers=headers, + ) + if response.status_code != 200: + return self.create_text_message("Failed to generate image") + + image = response.json()["images"][0] + + return self.create_blob_message( + blob=b64decode(image), + meta={"mime_type": "image/png"}, + save_as=self.VariableKey.IMAGE.value, + ) + + except Exception as e: + return self.create_text_message("Failed to generate image") + + def text2img( + self, base_url: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + generate image + """ + # copy draw options + draw_options = deepcopy(DRAW_TEXT_OPTIONS) + draw_options.update(tool_parameters) + # get prompt lora model + prompt = tool_parameters.get("prompt", "") + lora = tool_parameters.get("lora", "") + model = tool_parameters.get("model", "") + if lora: + draw_options["prompt"] = f"{lora},{prompt}" + else: + draw_options["prompt"] = prompt + draw_options["override_settings"]["sd_model_checkpoint"] = model + api_key = self.runtime.credentials.get("api_key") or "abc" + headers = {"Authorization": f"Bearer {api_key}"} + try: + url = str(URL(base_url) / "sdapi" / "v1" / "txt2img") + response = post( + url, + json=draw_options, + timeout=120, + headers=headers, + ) + if response.status_code != 200: + return self.create_text_message("Failed to generate image") + + image = response.json()["images"][0] + + return self.create_blob_message( + blob=b64decode(image), + meta={"mime_type": "image/png"}, + save_as=self.VariableKey.IMAGE.value, + ) + + except Exception as e: + return self.create_text_message("Failed to generate image") + + def get_runtime_parameters(self) -> list[ToolParameter]: + parameters = [ + ToolParameter( + name="prompt", + label=I18nObject(en_US="Prompt", zh_Hans="Prompt"), + human_description=I18nObject( + en_US="Image prompt, you can check the official documentation of Stable Diffusion", + zh_Hans="图像提示词,您可以查看 Stable Diffusion 的官方文档", + ), + type=ToolParameter.ToolParameterType.STRING, + form=ToolParameter.ToolParameterForm.LLM, + llm_description="Image prompt of Stable Diffusion, you should describe the image you want to generate" + " as a list of words as possible as detailed, the prompt must be written in English.", + required=True, + ), + ] + if len(self.list_default_image_variables()) != 0: + parameters.append( + ToolParameter( + name="image_id", + label=I18nObject(en_US="image_id", zh_Hans="image_id"), + human_description=I18nObject( + en_US="Image id of the image you want to generate based on, if you want to generate image based" + " on the default image, you can leave this field empty.", + zh_Hans="您想要生成的图像的图像 ID,如果您想要基于默认图像生成图像,则可以将此字段留空。", + ), + type=ToolParameter.ToolParameterType.STRING, + form=ToolParameter.ToolParameterForm.LLM, + llm_description="Image id of the original image, you can leave this field empty if you want to" + " generate a new image.", + required=True, + options=[ + ToolParameterOption(value=i.name, label=I18nObject(en_US=i.name, zh_Hans=i.name)) + for i in self.list_default_image_variables() + ], + ) + ) + + if self.runtime.credentials: + try: + models = self.get_sd_models() + if len(models) != 0: + parameters.append( + ToolParameter( + name="model", + label=I18nObject(en_US="Model", zh_Hans="Model"), + human_description=I18nObject( + en_US="Model of Stable Diffusion, you can check the official documentation" + " of Stable Diffusion", + zh_Hans="Stable Diffusion 的模型,您可以查看 Stable Diffusion 的官方文档", + ), + type=ToolParameter.ToolParameterType.SELECT, + form=ToolParameter.ToolParameterForm.FORM, + llm_description="Model of Stable Diffusion, you can check the official documentation" + " of Stable Diffusion", + required=True, + default=models[0], + options=[ + ToolParameterOption(value=i, label=I18nObject(en_US=i, zh_Hans=i)) for i in models + ], + ) + ) + + except: + pass + + sample_methods = self.get_sample_methods() + if len(sample_methods) != 0: + parameters.append( + ToolParameter( + name="sampler_name", + label=I18nObject(en_US="Sampling method", zh_Hans="Sampling method"), + human_description=I18nObject( + en_US="Sampling method of Stable Diffusion, you can check the official documentation" + " of Stable Diffusion", + zh_Hans="Stable Diffusion 的Sampling method,您可以查看 Stable Diffusion 的官方文档", + ), + type=ToolParameter.ToolParameterType.SELECT, + form=ToolParameter.ToolParameterForm.FORM, + llm_description="Sampling method of Stable Diffusion, you can check the official documentation" + " of Stable Diffusion", + required=True, + default=sample_methods[0], + options=[ + ToolParameterOption(value=i, label=I18nObject(en_US=i, zh_Hans=i)) for i in sample_methods + ], + ) + ) + return parameters diff --git a/api/core/tools/provider/builtin/xinference/tools/stable_diffusion.yaml b/api/core/tools/provider/builtin/xinference/tools/stable_diffusion.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4f1d17f175c5677b4585f2b442edea57f96db891 --- /dev/null +++ b/api/core/tools/provider/builtin/xinference/tools/stable_diffusion.yaml @@ -0,0 +1,87 @@ +identity: + name: stable_diffusion + author: xinference + label: + en_US: Stable Diffusion + zh_Hans: Stable Diffusion +description: + human: + en_US: Generate images using Stable Diffusion models. + zh_Hans: 使用 Stable Diffusion 模型生成图片。 + llm: draw the image you want based on your prompt. +parameters: + - name: prompt + type: string + required: true + label: + en_US: Prompt + zh_Hans: 提示词 + human_description: + en_US: Image prompt + zh_Hans: 图像提示词 + llm_description: Image prompt of Stable Diffusion, you should describe the image you want to generate as a list of words as possible as detailed, the prompt must be written in English. + form: llm + - name: model + type: string + required: false + label: + en_US: Model Name + zh_Hans: 模型名称 + human_description: + en_US: Model Name + zh_Hans: 模型名称 + form: form + - name: lora + type: string + required: false + label: + en_US: Lora + zh_Hans: Lora + human_description: + en_US: Lora + zh_Hans: Lora + form: form + - name: steps + type: number + required: false + label: + en_US: Steps + zh_Hans: Steps + human_description: + en_US: Steps + zh_Hans: Steps + form: form + default: 10 + - name: width + type: number + required: false + label: + en_US: Width + zh_Hans: Width + human_description: + en_US: Width + zh_Hans: Width + form: form + default: 1024 + - name: height + type: number + required: false + label: + en_US: Height + zh_Hans: Height + human_description: + en_US: Height + zh_Hans: Height + form: form + default: 1024 + - name: negative_prompt + type: string + required: false + label: + en_US: Negative prompt + zh_Hans: Negative prompt + human_description: + en_US: Negative prompt + zh_Hans: Negative prompt + form: form + default: bad art, ugly, deformed, watermark, duplicated, discontinuous lines diff --git a/api/core/tools/provider/builtin/xinference/xinference.py b/api/core/tools/provider/builtin/xinference/xinference.py new file mode 100644 index 0000000000000000000000000000000000000000..9692e4060e8a87b8cac45f7bf057fb4241a57583 --- /dev/null +++ b/api/core/tools/provider/builtin/xinference/xinference.py @@ -0,0 +1,24 @@ +import requests + +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class XinferenceProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + base_url = credentials.get("base_url", "").removesuffix("/") + api_key = credentials.get("api_key", "") + if not api_key: + api_key = "abc" + credentials["api_key"] = api_key + model = credentials.get("model", "") + if not base_url or not model: + raise ToolProviderCredentialValidationError("Xinference base_url and model is required") + headers = {"Authorization": f"Bearer {api_key}"} + res = requests.post( + f"{base_url}/sdapi/v1/options", + headers=headers, + json={"sd_model_checkpoint": model}, + ) + if res.status_code != 200: + raise ToolProviderCredentialValidationError("Xinference API key is invalid") diff --git a/api/core/tools/provider/builtin/xinference/xinference.yaml b/api/core/tools/provider/builtin/xinference/xinference.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b0c02b9cbcb01a7cf719265db7f78fbc9e0671ee --- /dev/null +++ b/api/core/tools/provider/builtin/xinference/xinference.yaml @@ -0,0 +1,40 @@ +identity: + author: xinference + name: xinference + label: + en_US: Xinference + zh_Hans: Xinference + description: + zh_Hans: Xinference 提供的兼容 Stable Diffusion web ui 的图片生成 API。 + en_US: Stable Diffusion web ui compatible API provided by Xinference. + icon: icon.png + tags: + - image +credentials_for_provider: + base_url: + type: secret-input + required: true + label: + en_US: Base URL + zh_Hans: Xinference 服务器的 Base URL + placeholder: + en_US: Please input Xinference server's Base URL + zh_Hans: 请输入 Xinference 服务器的 Base URL + model: + type: text-input + required: true + label: + en_US: Model + zh_Hans: 模型 + placeholder: + en_US: Please input your model name + zh_Hans: 请输入你的模型名称 + api_key: + type: secret-input + required: false + label: + en_US: API Key + zh_Hans: Xinference 服务器的 API Key + placeholder: + en_US: Please input Xinference server's API Key + zh_Hans: 请输入 Xinference 服务器的 API Key diff --git a/api/core/tools/provider/builtin/yahoo/_assets/icon.png b/api/core/tools/provider/builtin/yahoo/_assets/icon.png new file mode 100644 index 0000000000000000000000000000000000000000..35d756f75410dbdf74ca14c8fba6e660e20b27d8 Binary files /dev/null and b/api/core/tools/provider/builtin/yahoo/_assets/icon.png differ diff --git a/api/core/tools/provider/builtin/yahoo/tools/analytics.py b/api/core/tools/provider/builtin/yahoo/tools/analytics.py new file mode 100644 index 0000000000000000000000000000000000000000..95a65ba22fc8afc2a2f88aaa9d93b8c3d4d0f1df --- /dev/null +++ b/api/core/tools/provider/builtin/yahoo/tools/analytics.py @@ -0,0 +1,70 @@ +from datetime import datetime +from typing import Any, Union + +import pandas as pd +from requests.exceptions import HTTPError, ReadTimeout +from yfinance import download # type: ignore + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class YahooFinanceAnalyticsTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + symbol = tool_parameters.get("symbol", "") + if not symbol: + return self.create_text_message("Please input symbol") + + time_range = [None, None] + start_date = tool_parameters.get("start_date", "") + if start_date: + time_range[0] = start_date + else: + time_range[0] = "1800-01-01" + + end_date = tool_parameters.get("end_date", "") + if end_date: + time_range[1] = end_date + else: + time_range[1] = datetime.now().strftime("%Y-%m-%d") + + stock_data = download(symbol, start=time_range[0], end=time_range[1]) + max_segments = min(15, len(stock_data)) + rows_per_segment = len(stock_data) // (max_segments or 1) + summary_data = [] + for i in range(max_segments): + start_idx = i * rows_per_segment + end_idx = (i + 1) * rows_per_segment if i < max_segments - 1 else len(stock_data) + segment_data = stock_data.iloc[start_idx:end_idx] + segment_summary = { + "Start Date": segment_data.index[0], + "End Date": segment_data.index[-1], + "Average Close": segment_data["Close"].mean(), + "Average Volume": segment_data["Volume"].mean(), + "Average Open": segment_data["Open"].mean(), + "Average High": segment_data["High"].mean(), + "Average Low": segment_data["Low"].mean(), + "Average Adj Close": segment_data["Adj Close"].mean(), + "Max Close": segment_data["Close"].max(), + "Min Close": segment_data["Close"].min(), + "Max Volume": segment_data["Volume"].max(), + "Min Volume": segment_data["Volume"].min(), + "Max Open": segment_data["Open"].max(), + "Min Open": segment_data["Open"].min(), + "Max High": segment_data["High"].max(), + "Min High": segment_data["High"].min(), + } + + summary_data.append(segment_summary) + + summary_df = pd.DataFrame(summary_data) + + try: + return self.create_text_message(str(summary_df.to_dict())) + except (HTTPError, ReadTimeout): + return self.create_text_message("There is a internet connection problem. Please try again later.") diff --git a/api/core/tools/provider/builtin/yahoo/tools/analytics.yaml b/api/core/tools/provider/builtin/yahoo/tools/analytics.yaml new file mode 100644 index 0000000000000000000000000000000000000000..89e66fb58149089211041ed0744cee4cc3418bfd --- /dev/null +++ b/api/core/tools/provider/builtin/yahoo/tools/analytics.yaml @@ -0,0 +1,54 @@ +identity: + name: yahoo_finance_analytics + author: Dify + label: + en_US: Analytics + zh_Hans: 分析 + pt_BR: Análises + icon: icon.svg +description: + human: + en_US: A tool for get analytics about a ticker from Yahoo Finance. + zh_Hans: 一个用于从雅虎财经获取分析数据的工具。 + pt_BR: Uma ferramenta para obter análises sobre um ticker do Yahoo Finance. + llm: A tool for get analytics from Yahoo Finance. Input should be the ticker symbol like AAPL. +parameters: + - name: symbol + type: string + required: true + label: + en_US: Ticker symbol + zh_Hans: 股票代码 + pt_BR: Símbolo do ticker + human_description: + en_US: The ticker symbol of the company you want to analyze. + zh_Hans: 你想要搜索的公司的股票代码。 + pt_BR: O símbolo do ticker da empresa que você deseja analisar. + llm_description: The ticker symbol of the company you want to analyze. + form: llm + - name: start_date + type: string + required: false + label: + en_US: Start date + zh_Hans: 开始日期 + pt_BR: Data de início + human_description: + en_US: The start date of the analytics. + zh_Hans: 分析的开始日期。 + pt_BR: A data de início das análises. + llm_description: The start date of the analytics, the format of the date must be YYYY-MM-DD like 2020-01-01. + form: llm + - name: end_date + type: string + required: false + label: + en_US: End date + zh_Hans: 结束日期 + pt_BR: Data de término + human_description: + en_US: The end date of the analytics. + zh_Hans: 分析的结束日期。 + pt_BR: A data de término das análises. + llm_description: The end date of the analytics, the format of the date must be YYYY-MM-DD like 2024-01-01. + form: llm diff --git a/api/core/tools/provider/builtin/yahoo/tools/news.py b/api/core/tools/provider/builtin/yahoo/tools/news.py new file mode 100644 index 0000000000000000000000000000000000000000..c9ae0c4ca7fcc67c3241fea779d1db3dcd9099a6 --- /dev/null +++ b/api/core/tools/provider/builtin/yahoo/tools/news.py @@ -0,0 +1,46 @@ +from typing import Any, Union + +import yfinance # type: ignore +from requests.exceptions import HTTPError, ReadTimeout + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class YahooFinanceSearchTickerTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + + query = tool_parameters.get("symbol", "") + if not query: + return self.create_text_message("Please input symbol") + + try: + return self.run(ticker=query, user_id=user_id) + except (HTTPError, ReadTimeout): + return self.create_text_message("There is a internet connection problem. Please try again later.") + + def run(self, ticker: str, user_id: str) -> ToolInvokeMessage: + company = yfinance.Ticker(ticker) + try: + if company.isin is None: + return self.create_text_message(f"Company ticker {ticker} not found.") + except (HTTPError, ReadTimeout, ConnectionError): + return self.create_text_message(f"Company ticker {ticker} not found.") + + links = [] + try: + links = [n["link"] for n in company.news if n["type"] == "STORY"] + except (HTTPError, ReadTimeout, ConnectionError): + if not links: + return self.create_text_message(f"There is nothing about {ticker} ticker") + if not links: + return self.create_text_message(f"No news found for company that searched with {ticker} ticker.") + + result = "\n\n".join([self.get_url(link) for link in links]) + + return self.create_text_message(self.summary(user_id=user_id, content=result)) diff --git a/api/core/tools/provider/builtin/yahoo/tools/news.yaml b/api/core/tools/provider/builtin/yahoo/tools/news.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4118c1a82f280f768b73c632ba8e4f84bd870cf3 --- /dev/null +++ b/api/core/tools/provider/builtin/yahoo/tools/news.yaml @@ -0,0 +1,28 @@ +identity: + name: yahoo_finance_news + author: Dify + label: + en_US: News + zh_Hans: 新闻 + pt_BR: Notícias + icon: icon.svg +description: + human: + en_US: A tool for get news about a ticker from Yahoo Finance. + zh_Hans: 一个用于从雅虎财经获取新闻的工具。 + pt_BR: Uma ferramenta para obter notícias sobre um ticker da Yahoo Finance. + llm: A tool for get news from Yahoo Finance. Input should be the ticker symbol like AAPL. +parameters: + - name: symbol + type: string + required: true + label: + en_US: Ticker symbol + zh_Hans: 股票代码 + pt_BR: Símbolo do ticker + human_description: + en_US: The ticker symbol of the company you want to search. + zh_Hans: 你想要搜索的公司的股票代码。 + pt_BR: O símbolo do ticker da empresa que você deseja pesquisar. + llm_description: The ticker symbol of the company you want to search. + form: llm diff --git a/api/core/tools/provider/builtin/yahoo/tools/ticker.py b/api/core/tools/provider/builtin/yahoo/tools/ticker.py new file mode 100644 index 0000000000000000000000000000000000000000..74d0d25addf04b7e4407236392498860d60db68b --- /dev/null +++ b/api/core/tools/provider/builtin/yahoo/tools/ticker.py @@ -0,0 +1,27 @@ +from typing import Any, Union + +from requests.exceptions import HTTPError, ReadTimeout +from yfinance import Ticker # type: ignore + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class YahooFinanceSearchTickerTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + query = tool_parameters.get("symbol", "") + if not query: + return self.create_text_message("Please input symbol") + + try: + return self.create_text_message(self.run(ticker=query)) + except (HTTPError, ReadTimeout): + return self.create_text_message("There is a internet connection problem. Please try again later.") + + def run(self, ticker: str) -> str: + return str(Ticker(ticker).info) diff --git a/api/core/tools/provider/builtin/yahoo/tools/ticker.yaml b/api/core/tools/provider/builtin/yahoo/tools/ticker.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3c1ee9cf316be9d6bc4b84a1e58bc4f2cc5f2deb --- /dev/null +++ b/api/core/tools/provider/builtin/yahoo/tools/ticker.yaml @@ -0,0 +1,28 @@ +identity: + name: yahoo_finance_ticker + author: Dify + label: + en_US: Ticker + zh_Hans: 股票信息 + pt_BR: Ticker + icon: icon.svg +description: + human: + en_US: A tool for search ticker information from Yahoo Finance. + zh_Hans: 一个用于从雅虎财经搜索股票信息的工具。 + pt_BR: Uma ferramenta para buscar informações de ticker do Yahoo Finance. + llm: A tool for search ticker information from Yahoo Finance. Input should be the ticker symbol like AAPL. +parameters: + - name: symbol + type: string + required: true + label: + en_US: Ticker symbol + zh_Hans: 股票代码 + pt_BR: Símbolo do ticker + human_description: + en_US: The ticker symbol of the company you want to search. + zh_Hans: 你想要搜索的公司的股票代码。 + pt_BR: O símbolo do ticker da empresa que você deseja pesquisar. + llm_description: The ticker symbol of the company you want to search. + form: llm diff --git a/api/core/tools/provider/builtin/yahoo/yahoo.py b/api/core/tools/provider/builtin/yahoo/yahoo.py new file mode 100644 index 0000000000000000000000000000000000000000..8d82084e76970354efb1225c17af5fe48dc33d47 --- /dev/null +++ b/api/core/tools/provider/builtin/yahoo/yahoo.py @@ -0,0 +1,20 @@ +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.yahoo.tools.ticker import YahooFinanceSearchTickerTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class YahooFinanceProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + try: + YahooFinanceSearchTickerTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "ticker": "MSFT", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/yahoo/yahoo.yaml b/api/core/tools/provider/builtin/yahoo/yahoo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f1e82952c09ba45ac9d5ec820163bdb99e0fef35 --- /dev/null +++ b/api/core/tools/provider/builtin/yahoo/yahoo.yaml @@ -0,0 +1,16 @@ +identity: + author: Dify + name: yahoo + label: + en_US: YahooFinance + zh_Hans: 雅虎财经 + pt_BR: YahooFinance + description: + en_US: Finance, and Yahoo! get the latest news, stock quotes, and interactive chart with Yahoo! + zh_Hans: 雅虎财经,获取并整理出最新的新闻、股票报价等一切你想要的财经信息。 + pt_BR: Finance, and Yahoo! get the latest news, stock quotes, and interactive chart with Yahoo! + icon: icon.png + tags: + - business + - finance +credentials_for_provider: diff --git a/api/core/tools/provider/builtin/youtube/_assets/icon.svg b/api/core/tools/provider/builtin/youtube/_assets/icon.svg new file mode 100644 index 0000000000000000000000000000000000000000..83b0700fecbf30782d922a4e266946bbfd42dc83 --- /dev/null +++ b/api/core/tools/provider/builtin/youtube/_assets/icon.svg @@ -0,0 +1,11 @@ + + + + + + + + + \ No newline at end of file diff --git a/api/core/tools/provider/builtin/youtube/tools/videos.py b/api/core/tools/provider/builtin/youtube/tools/videos.py new file mode 100644 index 0000000000000000000000000000000000000000..a24fe89679b29bc311723a58ef50ba0fc2989aa4 --- /dev/null +++ b/api/core/tools/provider/builtin/youtube/tools/videos.py @@ -0,0 +1,74 @@ +from datetime import datetime +from typing import Any, Union + +from googleapiclient.discovery import build # type: ignore + +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool.builtin_tool import BuiltinTool + + +class YoutubeVideosAnalyticsTool(BuiltinTool): + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke tools + """ + channel = tool_parameters.get("channel", "") + if not channel: + return self.create_text_message("Please input symbol") + + time_range = [None, None] + start_date = tool_parameters.get("start_date", "") + if start_date: + time_range[0] = start_date + else: + time_range[0] = "1800-01-01" + + end_date = tool_parameters.get("end_date", "") + if end_date: + time_range[1] = end_date + else: + time_range[1] = datetime.now().strftime("%Y-%m-%d") + + if "google_api_key" not in self.runtime.credentials or not self.runtime.credentials["google_api_key"]: + return self.create_text_message("Please input api key") + + youtube = build("youtube", "v3", developerKey=self.runtime.credentials["google_api_key"]) + + # try to get channel id + search_results = youtube.search().list(q=channel, type="channel", order="relevance", part="id").execute() + channel_id = search_results["items"][0]["id"]["channelId"] + + start_date, end_date = time_range + + start_date = datetime.strptime(start_date, "%Y-%m-%d").strftime("%Y-%m-%dT%H:%M:%SZ") + end_date = datetime.strptime(end_date, "%Y-%m-%d").strftime("%Y-%m-%dT%H:%M:%SZ") + + # get videos + time_range_videos = ( + youtube.search() + .list( + part="snippet", + channelId=channel_id, + order="date", + type="video", + publishedAfter=start_date, + publishedBefore=end_date, + ) + .execute() + ) + + def extract_video_data(video_list): + data = [] + for video in video_list["items"]: + video_id = video["id"]["videoId"] + video_info = youtube.videos().list(part="snippet,statistics", id=video_id).execute() + title = video_info["items"][0]["snippet"]["title"] + views = video_info["items"][0]["statistics"]["viewCount"] + data.append({"Title": title, "Views": views}) + return data + + summary = extract_video_data(time_range_videos) + + return self.create_text_message(str(summary)) diff --git a/api/core/tools/provider/builtin/youtube/tools/videos.yaml b/api/core/tools/provider/builtin/youtube/tools/videos.yaml new file mode 100644 index 0000000000000000000000000000000000000000..976699eb6279106809c64cbc68094cd9b82af095 --- /dev/null +++ b/api/core/tools/provider/builtin/youtube/tools/videos.yaml @@ -0,0 +1,54 @@ +identity: + name: youtube_video_statistics + author: Dify + label: + en_US: Video statistics + zh_Hans: 视频统计 + pt_BR: Estatísticas de vídeo + icon: icon.svg +description: + human: + en_US: A tool for get statistics about a channel's videos. + zh_Hans: 一个用于获取油管频道视频统计数据的工具。 + pt_BR: Uma ferramenta para obter estatísticas sobre os vídeos de um canal. + llm: A tool for get statistics about a channel's videos. Input should be the name of the channel like PewDiePie. +parameters: + - name: channel + type: string + required: true + label: + en_US: Channel name + zh_Hans: 频道名 + pt_BR: Nome do canal + human_description: + en_US: The name of the channel you want to search. + zh_Hans: 你想要搜索的油管频道名。 + pt_BR: O nome do canal que você deseja pesquisar. + llm_description: The name of the channel you want to search. + form: llm + - name: start_date + type: string + required: false + label: + en_US: Start date + zh_Hans: 开始日期 + pt_BR: Data de início + human_description: + en_US: The start date of the analytics. + zh_Hans: 分析的开始日期。 + pt_BR: A data de início da análise. + llm_description: The start date of the analytics, the format of the date must be YYYY-MM-DD like 2020-01-01. + form: llm + - name: end_date + type: string + required: false + label: + en_US: End date + zh_Hans: 结束日期 + pt_BR: Data de término + human_description: + en_US: The end date of the analytics. + zh_Hans: 分析的结束日期。 + pt_BR: A data de término da análise. + llm_description: The end date of the analytics, the format of the date must be YYYY-MM-DD like 2024-01-01. + form: llm diff --git a/api/core/tools/provider/builtin/youtube/youtube.py b/api/core/tools/provider/builtin/youtube/youtube.py new file mode 100644 index 0000000000000000000000000000000000000000..07e430bcbf27e1789b895f6865c4f5d47fe812d2 --- /dev/null +++ b/api/core/tools/provider/builtin/youtube/youtube.py @@ -0,0 +1,22 @@ +from core.tools.errors import ToolProviderCredentialValidationError +from core.tools.provider.builtin.youtube.tools.videos import YoutubeVideosAnalyticsTool +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController + + +class YahooFinanceProvider(BuiltinToolProviderController): + def _validate_credentials(self, credentials: dict) -> None: + try: + YoutubeVideosAnalyticsTool().fork_tool_runtime( + runtime={ + "credentials": credentials, + } + ).invoke( + user_id="", + tool_parameters={ + "channel": "UC2JZCsZSOudXA08cMMRCL9g", + "start_date": "2020-01-01", + "end_date": "2024-12-31", + }, + ) + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) diff --git a/api/core/tools/provider/builtin/youtube/youtube.yaml b/api/core/tools/provider/builtin/youtube/youtube.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d6915b9a3247672a4d06ab627ecc43d526f2b6d1 --- /dev/null +++ b/api/core/tools/provider/builtin/youtube/youtube.yaml @@ -0,0 +1,31 @@ +identity: + author: Dify + name: youtube + label: + en_US: YouTube + zh_Hans: YouTube + pt_BR: YouTube + description: + en_US: YouTube + zh_Hans: YouTube(油管)是全球最大的视频分享网站,用户可以在上面上传、观看和分享视频。 + pt_BR: YouTube é o maior site de compartilhamento de vídeos do mundo, onde os usuários podem fazer upload, assistir e compartilhar vídeos. + icon: icon.svg + tags: + - videos +credentials_for_provider: + google_api_key: + type: secret-input + required: true + label: + en_US: Google API key + zh_Hans: Google API key + pt_BR: Chave da API do Google + placeholder: + en_US: Please input your Google API key + zh_Hans: 请输入你的 Google API key + pt_BR: Insira sua chave da API do Google + help: + en_US: Get your Google API key from Google + zh_Hans: 从 Google 获取您的 Google API key + pt_BR: Obtenha sua chave da API do Google no Google + url: https://console.developers.google.com/apis/credentials diff --git a/api/core/tools/provider/builtin_tool_provider.py b/api/core/tools/provider/builtin_tool_provider.py new file mode 100644 index 0000000000000000000000000000000000000000..61de75ac5e2ccd17d5861d5ee189c7b47dc85da4 --- /dev/null +++ b/api/core/tools/provider/builtin_tool_provider.py @@ -0,0 +1,244 @@ +from abc import abstractmethod +from os import listdir, path +from typing import Any, Optional + +from core.helper.module_import_helper import load_single_subclass_from_source +from core.tools.entities.tool_entities import ToolParameter, ToolProviderCredentials, ToolProviderType +from core.tools.entities.values import ToolLabelEnum, default_tool_label_dict +from core.tools.errors import ( + ToolNotFoundError, + ToolParameterValidationError, + ToolProviderNotFoundError, +) +from core.tools.provider.tool_provider import ToolProviderController +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.tool.tool import Tool +from core.tools.utils.yaml_utils import load_yaml_file + + +class BuiltinToolProviderController(ToolProviderController): + def __init__(self, **data: Any) -> None: + if self.provider_type in {ToolProviderType.API, ToolProviderType.APP}: + super().__init__(**data) + return + + # load provider yaml + provider = self.__class__.__module__.split(".")[-1] + yaml_path = path.join(path.dirname(path.realpath(__file__)), "builtin", provider, f"{provider}.yaml") + try: + provider_yaml = load_yaml_file(yaml_path, ignore_error=False) + except Exception as e: + raise ToolProviderNotFoundError(f"can not load provider yaml for {provider}: {e}") + + if "credentials_for_provider" in provider_yaml and provider_yaml["credentials_for_provider"] is not None: + # set credentials name + for credential_name in provider_yaml["credentials_for_provider"]: + provider_yaml["credentials_for_provider"][credential_name]["name"] = credential_name + + super().__init__( + **{ + "identity": provider_yaml["identity"], + "credentials_schema": provider_yaml.get("credentials_for_provider", None), + } + ) + + def _get_builtin_tools(self) -> list[Tool]: + """ + returns a list of tools that the provider can provide + + :return: list of tools + """ + if self.tools: + return self.tools + if not self.identity: + return [] + + provider = self.identity.name + tool_path = path.join(path.dirname(path.realpath(__file__)), "builtin", provider, "tools") + # get all the yaml files in the tool path + tool_files = list(filter(lambda x: x.endswith(".yaml") and not x.startswith("__"), listdir(tool_path))) + tools = [] + for tool_file in tool_files: + # get tool name + tool_name = tool_file.split(".")[0] + tool = load_yaml_file(path.join(tool_path, tool_file), ignore_error=False) + + # get tool class, import the module + assistant_tool_class = load_single_subclass_from_source( + module_name=f"core.tools.provider.builtin.{provider}.tools.{tool_name}", + script_path=path.join( + path.dirname(path.realpath(__file__)), "builtin", provider, "tools", f"{tool_name}.py" + ), + parent_type=BuiltinTool, + ) + tool["identity"]["provider"] = provider + tools.append(assistant_tool_class(**tool)) + + self.tools = tools + return tools + + def get_credentials_schema(self) -> dict[str, ToolProviderCredentials]: + """ + returns the credentials schema of the provider + + :return: the credentials schema + """ + if not self.credentials_schema: + return {} + + return self.credentials_schema.copy() + + def get_tools(self, user_id: str = "", tenant_id: str = "") -> Optional[list[Tool]]: + """ + returns a list of tools that the provider can provide + + :return: list of tools + """ + return self._get_builtin_tools() + + def get_tool(self, tool_name: str) -> Optional[Tool]: + """ + returns the tool that the provider can provide + """ + tools = self.get_tools() + if tools is None: + raise ValueError("tools not found") + return next((t for t in tools if t.identity and t.identity.name == tool_name), None) + + def get_parameters(self, tool_name: str) -> list[ToolParameter]: + """ + returns the parameters of the tool + + :param tool_name: the name of the tool, defined in `get_tools` + :return: list of parameters + """ + tools = self.get_tools() + if tools is None: + raise ToolNotFoundError(f"tool {tool_name} not found") + tool = next((t for t in tools if t.identity and t.identity.name == tool_name), None) + if tool is None: + raise ToolNotFoundError(f"tool {tool_name} not found") + return tool.parameters or [] + + @property + def need_credentials(self) -> bool: + """ + returns whether the provider needs credentials + + :return: whether the provider needs credentials + """ + return self.credentials_schema is not None and len(self.credentials_schema) != 0 + + @property + def provider_type(self) -> ToolProviderType: + """ + returns the type of the provider + + :return: type of the provider + """ + return ToolProviderType.BUILT_IN + + @property + def tool_labels(self) -> list[str]: + """ + returns the labels of the provider + + :return: labels of the provider + """ + label_enums = self._get_tool_labels() + return [default_tool_label_dict[label].name for label in label_enums] + + def _get_tool_labels(self) -> list[ToolLabelEnum]: + """ + returns the labels of the provider + """ + if self.identity is None: + return [] + return self.identity.tags or [] + + def validate_parameters(self, tool_id: int, tool_name: str, tool_parameters: dict[str, Any]) -> None: + """ + validate the parameters of the tool and set the default value if needed + + :param tool_name: the name of the tool, defined in `get_tools` + :param tool_parameters: the parameters of the tool + """ + tool_parameters_schema = self.get_parameters(tool_name) + + tool_parameters_need_to_validate: dict[str, ToolParameter] = {} + for parameter in tool_parameters_schema: + tool_parameters_need_to_validate[parameter.name] = parameter + + for parameter_name in tool_parameters: + if parameter_name not in tool_parameters_need_to_validate: + raise ToolParameterValidationError(f"parameter {parameter_name} not found in tool {tool_name}") + + # check type + parameter_schema = tool_parameters_need_to_validate[parameter_name] + if parameter_schema.type == ToolParameter.ToolParameterType.STRING: + if not isinstance(tool_parameters[parameter_name], str): + raise ToolParameterValidationError(f"parameter {parameter_name} should be string") + + elif parameter_schema.type == ToolParameter.ToolParameterType.NUMBER: + if not isinstance(tool_parameters[parameter_name], int | float): + raise ToolParameterValidationError(f"parameter {parameter_name} should be number") + + if parameter_schema.min is not None and tool_parameters[parameter_name] < parameter_schema.min: + raise ToolParameterValidationError( + f"parameter {parameter_name} should be greater than {parameter_schema.min}" + ) + + if parameter_schema.max is not None and tool_parameters[parameter_name] > parameter_schema.max: + raise ToolParameterValidationError( + f"parameter {parameter_name} should be less than {parameter_schema.max}" + ) + + elif parameter_schema.type == ToolParameter.ToolParameterType.BOOLEAN: + if not isinstance(tool_parameters[parameter_name], bool): + raise ToolParameterValidationError(f"parameter {parameter_name} should be boolean") + + elif parameter_schema.type == ToolParameter.ToolParameterType.SELECT: + if not isinstance(tool_parameters[parameter_name], str): + raise ToolParameterValidationError(f"parameter {parameter_name} should be string") + + options = parameter_schema.options + if not isinstance(options, list): + raise ToolParameterValidationError(f"parameter {parameter_name} options should be list") + + if tool_parameters[parameter_name] not in [x.value for x in options]: + raise ToolParameterValidationError(f"parameter {parameter_name} should be one of {options}") + + tool_parameters_need_to_validate.pop(parameter_name) + + for parameter_name in tool_parameters_need_to_validate: + parameter_schema = tool_parameters_need_to_validate[parameter_name] + if parameter_schema.required: + raise ToolParameterValidationError(f"parameter {parameter_name} is required") + + # the parameter is not set currently, set the default value if needed + if parameter_schema.default is not None: + default_value = parameter_schema.type.cast_value(parameter_schema.default) + tool_parameters[parameter_name] = default_value + + def validate_credentials(self, credentials: dict[str, Any]) -> None: + """ + validate the credentials of the provider + + :param tool_name: the name of the tool, defined in `get_tools` + :param credentials: the credentials of the tool + """ + # validate credentials format + self.validate_credentials_format(credentials) + + # validate credentials + self._validate_credentials(credentials) + + @abstractmethod + def _validate_credentials(self, credentials: dict[str, Any]) -> None: + """ + validate the credentials of the provider + + :param tool_name: the name of the tool, defined in `get_tools` + :param credentials: the credentials of the tool + """ + pass diff --git a/api/core/tools/provider/tool_provider.py b/api/core/tools/provider/tool_provider.py new file mode 100644 index 0000000000000000000000000000000000000000..e35207e4f06404c2ef6b7ca359cc05885fe46b07 --- /dev/null +++ b/api/core/tools/provider/tool_provider.py @@ -0,0 +1,201 @@ +from abc import ABC, abstractmethod +from typing import Any, Optional + +from pydantic import BaseModel + +from core.tools.entities.tool_entities import ( + ToolParameter, + ToolProviderCredentials, + ToolProviderIdentity, + ToolProviderType, +) +from core.tools.errors import ToolNotFoundError, ToolParameterValidationError, ToolProviderCredentialValidationError +from core.tools.tool.tool import Tool + + +class ToolProviderController(BaseModel, ABC): + identity: Optional[ToolProviderIdentity] = None + tools: Optional[list[Tool]] = None + credentials_schema: Optional[dict[str, ToolProviderCredentials]] = None + + def get_credentials_schema(self) -> dict[str, ToolProviderCredentials]: + """ + returns the credentials schema of the provider + + :return: the credentials schema + """ + if self.credentials_schema is None: + return {} + return self.credentials_schema.copy() + + @abstractmethod + def get_tools(self, user_id: str = "", tenant_id: str = "") -> Optional[list[Tool]]: + """ + returns a list of tools that the provider can provide + + :return: list of tools + """ + pass + + @abstractmethod + def get_tool(self, tool_name: str) -> Optional[Tool]: + """ + returns a tool that the provider can provide + + :return: tool + """ + pass + + def get_parameters(self, tool_name: str) -> list[ToolParameter]: + """ + returns the parameters of the tool + + :param tool_name: the name of the tool, defined in `get_tools` + :return: list of parameters + """ + tools = self.get_tools() + if tools is None: + raise ToolNotFoundError(f"tool {tool_name} not found") + tool = next((t for t in tools if t.identity and t.identity.name == tool_name), None) + if tool is None: + raise ToolNotFoundError(f"tool {tool_name} not found") + return tool.parameters or [] + + @property + def provider_type(self) -> ToolProviderType: + """ + returns the type of the provider + + :return: type of the provider + """ + return ToolProviderType.BUILT_IN + + def validate_parameters(self, tool_id: int, tool_name: str, tool_parameters: dict[str, Any]) -> None: + """ + validate the parameters of the tool and set the default value if needed + + :param tool_name: the name of the tool, defined in `get_tools` + :param tool_parameters: the parameters of the tool + """ + tool_parameters_schema = self.get_parameters(tool_name) + + tool_parameters_need_to_validate: dict[str, ToolParameter] = {} + for parameter in tool_parameters_schema: + tool_parameters_need_to_validate[parameter.name] = parameter + + for tool_parameter in tool_parameters: + if tool_parameter not in tool_parameters_need_to_validate: + raise ToolParameterValidationError(f"parameter {tool_parameter} not found in tool {tool_name}") + + # check type + parameter_schema = tool_parameters_need_to_validate[tool_parameter] + if parameter_schema.type == ToolParameter.ToolParameterType.STRING: + if not isinstance(tool_parameters[tool_parameter], str): + raise ToolParameterValidationError(f"parameter {tool_parameter} should be string") + + elif parameter_schema.type == ToolParameter.ToolParameterType.NUMBER: + if not isinstance(tool_parameters[tool_parameter], int | float): + raise ToolParameterValidationError(f"parameter {tool_parameter} should be number") + + if parameter_schema.min is not None and tool_parameters[tool_parameter] < parameter_schema.min: + raise ToolParameterValidationError( + f"parameter {tool_parameter} should be greater than {parameter_schema.min}" + ) + + if parameter_schema.max is not None and tool_parameters[tool_parameter] > parameter_schema.max: + raise ToolParameterValidationError( + f"parameter {tool_parameter} should be less than {parameter_schema.max}" + ) + + elif parameter_schema.type == ToolParameter.ToolParameterType.BOOLEAN: + if not isinstance(tool_parameters[tool_parameter], bool): + raise ToolParameterValidationError(f"parameter {tool_parameter} should be boolean") + + elif parameter_schema.type == ToolParameter.ToolParameterType.SELECT: + if not isinstance(tool_parameters[tool_parameter], str): + raise ToolParameterValidationError(f"parameter {tool_parameter} should be string") + + options = parameter_schema.options + if not isinstance(options, list): + raise ToolParameterValidationError(f"parameter {tool_parameter} options should be list") + + if tool_parameters[tool_parameter] not in [x.value for x in options]: + raise ToolParameterValidationError(f"parameter {tool_parameter} should be one of {options}") + + tool_parameters_need_to_validate.pop(tool_parameter) + + for tool_parameter_validate in tool_parameters_need_to_validate: + parameter_schema = tool_parameters_need_to_validate[tool_parameter_validate] + if parameter_schema.required: + raise ToolParameterValidationError(f"parameter {tool_parameter_validate} is required") + + # the parameter is not set currently, set the default value if needed + if parameter_schema.default is not None: + tool_parameters[tool_parameter_validate] = parameter_schema.type.cast_value(parameter_schema.default) + + def validate_credentials_format(self, credentials: dict[str, Any]) -> None: + """ + validate the format of the credentials of the provider and set the default value if needed + + :param credentials: the credentials of the tool + """ + credentials_schema = self.credentials_schema + if credentials_schema is None: + return + + credentials_need_to_validate: dict[str, ToolProviderCredentials] = {} + for credential_name in credentials_schema: + credentials_need_to_validate[credential_name] = credentials_schema[credential_name] + + for credential_name in credentials: + if credential_name not in credentials_need_to_validate: + if self.identity is None: + raise ValueError("identity is not set") + raise ToolProviderCredentialValidationError( + f"credential {credential_name} not found in provider {self.identity.name}" + ) + + # check type + credential_schema = credentials_need_to_validate[credential_name] + if not credential_schema.required and credentials[credential_name] is None: + continue + + if credential_schema.type in { + ToolProviderCredentials.CredentialsType.SECRET_INPUT, + ToolProviderCredentials.CredentialsType.TEXT_INPUT, + }: + if not isinstance(credentials[credential_name], str): + raise ToolProviderCredentialValidationError(f"credential {credential_name} should be string") + + elif credential_schema.type == ToolProviderCredentials.CredentialsType.SELECT: + if not isinstance(credentials[credential_name], str): + raise ToolProviderCredentialValidationError(f"credential {credential_name} should be string") + + options = credential_schema.options + if not isinstance(options, list): + raise ToolProviderCredentialValidationError(f"credential {credential_name} options should be list") + + if credentials[credential_name] not in [x.value for x in options]: + raise ToolProviderCredentialValidationError( + f"credential {credential_name} should be one of {options}" + ) + + credentials_need_to_validate.pop(credential_name) + + for credential_name in credentials_need_to_validate: + credential_schema = credentials_need_to_validate[credential_name] + if credential_schema.required: + raise ToolProviderCredentialValidationError(f"credential {credential_name} is required") + + # the credential is not set currently, set the default value if needed + if credential_schema.default is not None: + default_value = credential_schema.default + # parse default value into the correct type + if credential_schema.type in { + ToolProviderCredentials.CredentialsType.SECRET_INPUT, + ToolProviderCredentials.CredentialsType.TEXT_INPUT, + ToolProviderCredentials.CredentialsType.SELECT, + }: + default_value = str(default_value) + + credentials[credential_name] = default_value diff --git a/api/core/tools/provider/workflow_tool_provider.py b/api/core/tools/provider/workflow_tool_provider.py new file mode 100644 index 0000000000000000000000000000000000000000..17fe2e20cf282e8affc06ab146dc48959639ded7 --- /dev/null +++ b/api/core/tools/provider/workflow_tool_provider.py @@ -0,0 +1,208 @@ +from typing import Optional + +from core.app.app_config.entities import VariableEntityType +from core.app.apps.workflow.app_config_manager import WorkflowAppConfigManager +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_entities import ( + ToolDescription, + ToolIdentity, + ToolParameter, + ToolParameterOption, + ToolProviderType, +) +from core.tools.provider.tool_provider import ToolProviderController +from core.tools.tool.tool import Tool +from core.tools.tool.workflow_tool import WorkflowTool +from core.tools.utils.workflow_configuration_sync import WorkflowToolConfigurationUtils +from extensions.ext_database import db +from models.model import App, AppMode +from models.tools import WorkflowToolProvider +from models.workflow import Workflow + +VARIABLE_TO_PARAMETER_TYPE_MAPPING = { + VariableEntityType.TEXT_INPUT: ToolParameter.ToolParameterType.STRING, + VariableEntityType.PARAGRAPH: ToolParameter.ToolParameterType.STRING, + VariableEntityType.SELECT: ToolParameter.ToolParameterType.SELECT, + VariableEntityType.NUMBER: ToolParameter.ToolParameterType.NUMBER, + VariableEntityType.FILE: ToolParameter.ToolParameterType.FILE, + VariableEntityType.FILE_LIST: ToolParameter.ToolParameterType.FILES, +} + + +class WorkflowToolProviderController(ToolProviderController): + provider_id: str + + @classmethod + def from_db(cls, db_provider: WorkflowToolProvider) -> "WorkflowToolProviderController": + app = db_provider.app + + if not app: + raise ValueError("app not found") + + controller = WorkflowToolProviderController.model_validate( + { + "identity": { + "author": db_provider.user.name if db_provider.user_id and db_provider.user else "", + "name": db_provider.label, + "label": {"en_US": db_provider.label, "zh_Hans": db_provider.label}, + "description": {"en_US": db_provider.description, "zh_Hans": db_provider.description}, + "icon": db_provider.icon, + }, + "credentials_schema": {}, + "provider_id": db_provider.id or "", + } + ) + + # init tools + + controller.tools = [controller._get_db_provider_tool(db_provider, app)] + + return controller + + @property + def provider_type(self) -> ToolProviderType: + return ToolProviderType.WORKFLOW + + def _get_db_provider_tool(self, db_provider: WorkflowToolProvider, app: App) -> WorkflowTool: + """ + get db provider tool + :param db_provider: the db provider + :param app: the app + :return: the tool + """ + workflow = ( + db.session.query(Workflow) + .filter(Workflow.app_id == db_provider.app_id, Workflow.version == db_provider.version) + .first() + ) + if not workflow: + raise ValueError("workflow not found") + + # fetch start node + graph = workflow.graph_dict + features_dict = workflow.features_dict + features = WorkflowAppConfigManager.convert_features(config_dict=features_dict, app_mode=AppMode.WORKFLOW) + + parameters = db_provider.parameter_configurations + variables = WorkflowToolConfigurationUtils.get_workflow_graph_variables(graph) + + def fetch_workflow_variable(variable_name: str): + return next(filter(lambda x: x.variable == variable_name, variables), None) + + user = db_provider.user + + workflow_tool_parameters = [] + for parameter in parameters: + variable = fetch_workflow_variable(parameter.name) + if variable: + parameter_type = None + options = None + if variable.type not in VARIABLE_TO_PARAMETER_TYPE_MAPPING: + raise ValueError(f"unsupported variable type {variable.type}") + parameter_type = VARIABLE_TO_PARAMETER_TYPE_MAPPING[variable.type] + + if variable.type == VariableEntityType.SELECT and variable.options: + options = [ + ToolParameterOption(value=option, label=I18nObject(en_US=option, zh_Hans=option)) + for option in variable.options + ] + + workflow_tool_parameters.append( + ToolParameter( + name=parameter.name, + label=I18nObject(en_US=variable.label, zh_Hans=variable.label), + human_description=I18nObject(en_US=parameter.description, zh_Hans=parameter.description), + type=parameter_type, + form=parameter.form, + llm_description=parameter.description, + required=variable.required, + options=options, + placeholder=I18nObject(en_US="", zh_Hans=""), + ) + ) + elif features.file_upload: + workflow_tool_parameters.append( + ToolParameter( + name=parameter.name, + label=I18nObject(en_US=parameter.name, zh_Hans=parameter.name), + human_description=I18nObject(en_US=parameter.description, zh_Hans=parameter.description), + type=ToolParameter.ToolParameterType.SYSTEM_FILES, + llm_description=parameter.description, + required=False, + form=parameter.form, + placeholder=I18nObject(en_US="", zh_Hans=""), + ) + ) + else: + raise ValueError("variable not found") + + return WorkflowTool( + identity=ToolIdentity( + author=user.name if user else "", + name=db_provider.name, + label=I18nObject(en_US=db_provider.label, zh_Hans=db_provider.label), + provider=self.provider_id, + icon=db_provider.icon, + ), + description=ToolDescription( + human=I18nObject(en_US=db_provider.description, zh_Hans=db_provider.description), + llm=db_provider.description, + ), + parameters=workflow_tool_parameters, + is_team_authorization=True, + workflow_app_id=app.id, + workflow_entities={ + "app": app, + "workflow": workflow, + }, + version=db_provider.version, + workflow_call_depth=0, + label=db_provider.label, + ) + + def get_tools(self, user_id: str = "", tenant_id: str = "") -> Optional[list[Tool]]: + """ + fetch tools from database + + :param user_id: the user id + :param tenant_id: the tenant id + :return: the tools + """ + if self.tools is not None: + return self.tools + + db_providers: Optional[WorkflowToolProvider] = ( + db.session.query(WorkflowToolProvider) + .filter( + WorkflowToolProvider.tenant_id == tenant_id, + WorkflowToolProvider.app_id == self.provider_id, + ) + .first() + ) + + if not db_providers: + return [] + if not db_providers.app: + raise ValueError("app not found") + + self.tools = [self._get_db_provider_tool(db_providers, db_providers.app)] + + return self.tools + + def get_tool(self, tool_name: str) -> Optional[Tool]: + """ + get tool by name + + :param tool_name: the name of the tool + :return: the tool + """ + if self.tools is None: + return None + + for tool in self.tools: + if tool.identity is None: + continue + if tool.identity.name == tool_name: + return tool + + return None diff --git a/api/core/tools/tool/api_tool.py b/api/core/tools/tool/api_tool.py new file mode 100644 index 0000000000000000000000000000000000000000..6904fecb46698ba8879436a706d9519e7a01678f --- /dev/null +++ b/api/core/tools/tool/api_tool.py @@ -0,0 +1,322 @@ +import json +from os import getenv +from typing import Any +from urllib.parse import urlencode + +import httpx + +from core.file.file_manager import download +from core.helper import ssrf_proxy +from core.tools.entities.tool_bundle import ApiToolBundle +from core.tools.entities.tool_entities import ToolInvokeMessage, ToolProviderType +from core.tools.errors import ToolInvokeError, ToolParameterValidationError, ToolProviderCredentialValidationError +from core.tools.tool.tool import Tool + +API_TOOL_DEFAULT_TIMEOUT = ( + int(getenv("API_TOOL_DEFAULT_CONNECT_TIMEOUT", "10")), + int(getenv("API_TOOL_DEFAULT_READ_TIMEOUT", "60")), +) + + +class ApiTool(Tool): + api_bundle: ApiToolBundle + + """ + Api tool + """ + + def fork_tool_runtime(self, runtime: dict[str, Any]) -> "Tool": + """ + fork a new tool with meta data + + :param meta: the meta data of a tool call processing, tenant_id is required + :return: the new tool + """ + if self.api_bundle is None: + raise ValueError("api_bundle is required") + return self.__class__( + identity=self.identity.model_copy() if self.identity else None, + parameters=self.parameters.copy() if self.parameters else None, + description=self.description.model_copy() if self.description else None, + api_bundle=self.api_bundle.model_copy(), + runtime=Tool.Runtime(**runtime), + ) + + def validate_credentials( + self, credentials: dict[str, Any], parameters: dict[str, Any], format_only: bool = False + ) -> str: + """ + validate the credentials for Api tool + """ + # assemble validate request and request parameters + headers = self.assembling_request(parameters) + + if format_only: + return "" + + response = self.do_http_request(self.api_bundle.server_url, self.api_bundle.method, headers, parameters) + # validate response + return self.validate_and_parse_response(response) + + def tool_provider_type(self) -> ToolProviderType: + return ToolProviderType.API + + def assembling_request(self, parameters: dict[str, Any]) -> dict[str, Any]: + headers = {} + if self.runtime is None: + raise ValueError("runtime is required") + credentials = self.runtime.credentials or {} + + if "auth_type" not in credentials: + raise ToolProviderCredentialValidationError("Missing auth_type") + + if credentials["auth_type"] == "api_key": + api_key_header = "api_key" + + if "api_key_header" in credentials: + api_key_header = credentials["api_key_header"] + + if "api_key_value" not in credentials: + raise ToolProviderCredentialValidationError("Missing api_key_value") + elif not isinstance(credentials["api_key_value"], str): + raise ToolProviderCredentialValidationError("api_key_value must be a string") + + if "api_key_header_prefix" in credentials: + api_key_header_prefix = credentials["api_key_header_prefix"] + if api_key_header_prefix == "basic" and credentials["api_key_value"]: + credentials["api_key_value"] = f"Basic {credentials['api_key_value']}" + elif api_key_header_prefix == "bearer" and credentials["api_key_value"]: + credentials["api_key_value"] = f"Bearer {credentials['api_key_value']}" + elif api_key_header_prefix == "custom": + pass + + headers[api_key_header] = credentials["api_key_value"] + + needed_parameters = [parameter for parameter in (self.api_bundle.parameters or []) if parameter.required] + for parameter in needed_parameters: + if parameter.required and parameter.name not in parameters: + raise ToolParameterValidationError(f"Missing required parameter {parameter.name}") + + if parameter.default is not None and parameter.name not in parameters: + parameters[parameter.name] = parameter.default + + return headers + + def validate_and_parse_response(self, response: httpx.Response) -> str: + """ + validate the response + """ + if isinstance(response, httpx.Response): + if response.status_code >= 400: + raise ToolInvokeError(f"Request failed with status code {response.status_code} and {response.text}") + if not response.content: + return "Empty response from the tool, please check your parameters and try again." + try: + response = response.json() + try: + return json.dumps(response, ensure_ascii=False) + except Exception as e: + return json.dumps(response) + except Exception as e: + return response.text + else: + raise ValueError(f"Invalid response type {type(response)}") + + @staticmethod + def get_parameter_value(parameter, parameters): + if parameter["name"] in parameters: + return parameters[parameter["name"]] + elif parameter.get("required", False): + raise ToolParameterValidationError(f"Missing required parameter {parameter['name']}") + else: + return (parameter.get("schema", {}) or {}).get("default", "") + + def do_http_request( + self, url: str, method: str, headers: dict[str, Any], parameters: dict[str, Any] + ) -> httpx.Response: + """ + do http request depending on api bundle + """ + method = method.lower() + + params = {} + path_params = {} + # FIXME: body should be a dict[str, Any] but it changed a lot in this function + body: Any = {} + cookies = {} + files = [] + + # check parameters + for parameter in self.api_bundle.openapi.get("parameters", []): + value = self.get_parameter_value(parameter, parameters) + if parameter["in"] == "path": + path_params[parameter["name"]] = value + + elif parameter["in"] == "query": + if value != "": + params[parameter["name"]] = value + + elif parameter["in"] == "cookie": + cookies[parameter["name"]] = value + + elif parameter["in"] == "header": + headers[parameter["name"]] = value + + # check if there is a request body and handle it + if "requestBody" in self.api_bundle.openapi and self.api_bundle.openapi["requestBody"] is not None: + # handle json request body + if "content" in self.api_bundle.openapi["requestBody"]: + for content_type in self.api_bundle.openapi["requestBody"]["content"]: + headers["Content-Type"] = content_type + body_schema = self.api_bundle.openapi["requestBody"]["content"][content_type]["schema"] + required = body_schema.get("required", []) + properties = body_schema.get("properties", {}) + for name, property in properties.items(): + if name in parameters: + if property.get("format") == "binary": + f = parameters[name] + files.append((name, (f.filename, download(f), f.mime_type))) + else: + # convert type + body[name] = self._convert_body_property_type(property, parameters[name]) + elif name in required: + raise ToolParameterValidationError( + f"Missing required parameter {name} in operation {self.api_bundle.operation_id}" + ) + elif "default" in property: + body[name] = property["default"] + else: + body[name] = None + break + + # replace path parameters + for name, value in path_params.items(): + url = url.replace(f"{{{name}}}", f"{value}") + + # parse http body data if needed + if "Content-Type" in headers: + if headers["Content-Type"] == "application/json": + body = json.dumps(body) + elif headers["Content-Type"] == "application/x-www-form-urlencoded": + body = urlencode(body) + else: + body = body + + if method in { + "get", + "head", + "post", + "put", + "delete", + "patch", + "options", + "GET", + "POST", + "PUT", + "PATCH", + "DELETE", + "HEAD", + "OPTIONS", + }: + response: httpx.Response = getattr(ssrf_proxy, method.lower())( + url, + params=params, + headers=headers, + cookies=cookies, + data=body, + files=files, + timeout=API_TOOL_DEFAULT_TIMEOUT, + follow_redirects=True, + ) + return response + else: + raise ValueError(f"Invalid http method {method}") + + def _convert_body_property_any_of( + self, property: dict[str, Any], value: Any, any_of: list[dict[str, Any]], max_recursive=10 + ) -> Any: + if max_recursive <= 0: + raise Exception("Max recursion depth reached") + for option in any_of or []: + try: + if "type" in option: + # Attempt to convert the value based on the type. + if option["type"] == "integer" or option["type"] == "int": + return int(value) + elif option["type"] == "number": + if "." in str(value): + return float(value) + else: + return int(value) + elif option["type"] == "string": + return str(value) + elif option["type"] == "boolean": + if str(value).lower() in {"true", "1"}: + return True + elif str(value).lower() in {"false", "0"}: + return False + else: + continue # Not a boolean, try next option + elif option["type"] == "null" and not value: + return None + else: + continue # Unsupported type, try next option + elif "anyOf" in option and isinstance(option["anyOf"], list): + # Recursive call to handle nested anyOf + return self._convert_body_property_any_of(property, value, option["anyOf"], max_recursive - 1) + except ValueError: + continue # Conversion failed, try next option + # If no option succeeded, you might want to return the value as is or raise an error + return value # or raise ValueError(f"Cannot convert value '{value}' to any specified type in anyOf") + + def _convert_body_property_type(self, property: dict[str, Any], value: Any) -> Any: + try: + if "type" in property: + if property["type"] == "integer" or property["type"] == "int": + return int(value) + elif property["type"] == "number": + # check if it is a float + if "." in str(value): + return float(value) + else: + return int(value) + elif property["type"] == "string": + return str(value) + elif property["type"] == "boolean": + return bool(value) + elif property["type"] == "null": + if value is None: + return None + elif property["type"] == "object" or property["type"] == "array": + if isinstance(value, str): + try: + return json.loads(value) + except ValueError: + return value + elif isinstance(value, dict): + return value + else: + return value + else: + raise ValueError(f"Invalid type {property['type']} for property {property}") + elif "anyOf" in property and isinstance(property["anyOf"], list): + return self._convert_body_property_any_of(property, value, property["anyOf"]) + except ValueError as e: + return value + + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage | list[ToolInvokeMessage]: + """ + invoke http request + """ + response: httpx.Response | str = "" + # assemble request + headers = self.assembling_request(tool_parameters) + + # do http request + response = self.do_http_request(self.api_bundle.server_url, self.api_bundle.method, headers, tool_parameters) + + # validate response + response = self.validate_and_parse_response(response) + + # assemble invoke message + return self.create_text_message(response) diff --git a/api/core/tools/tool/builtin_tool.py b/api/core/tools/tool/builtin_tool.py new file mode 100644 index 0000000000000000000000000000000000000000..adda4297f38e8a68bf9997a78f785aef8c53591d --- /dev/null +++ b/api/core/tools/tool/builtin_tool.py @@ -0,0 +1,144 @@ +from typing import Optional, cast + +from core.model_runtime.entities.llm_entities import LLMResult +from core.model_runtime.entities.message_entities import PromptMessage, SystemPromptMessage, UserPromptMessage +from core.tools.entities.tool_entities import ToolProviderType +from core.tools.tool.tool import Tool +from core.tools.utils.model_invocation_utils import ModelInvocationUtils +from core.tools.utils.web_reader_tool import get_url + +_SUMMARY_PROMPT = """You are a professional language researcher, you are interested in the language +and you can quickly aimed at the main point of an webpage and reproduce it in your own words but +retain the original meaning and keep the key points. +however, the text you got is too long, what you got is possible a part of the text. +Please summarize the text you got. +""" + + +class BuiltinTool(Tool): + """ + Builtin tool + + :param meta: the meta data of a tool call processing + """ + + def invoke_model(self, user_id: str, prompt_messages: list[PromptMessage], stop: list[str]) -> LLMResult: + """ + invoke model + + :param model_config: the model config + :param prompt_messages: the prompt messages + :param stop: the stop words + :return: the model result + """ + # invoke model + if self.runtime is None or self.identity is None: + raise ValueError("runtime and identity are required") + + return ModelInvocationUtils.invoke( + user_id=user_id, + tenant_id=self.runtime.tenant_id or "", + tool_type="builtin", + tool_name=self.identity.name, + prompt_messages=prompt_messages, + ) + + def tool_provider_type(self) -> ToolProviderType: + return ToolProviderType.BUILT_IN + + def get_max_tokens(self) -> int: + """ + get max tokens + + :param model_config: the model config + :return: the max tokens + """ + if self.runtime is None: + raise ValueError("runtime is required") + + return ModelInvocationUtils.get_max_llm_context_tokens( + tenant_id=self.runtime.tenant_id or "", + ) + + def get_prompt_tokens(self, prompt_messages: list[PromptMessage]) -> int: + """ + get prompt tokens + + :param prompt_messages: the prompt messages + :return: the tokens + """ + if self.runtime is None: + raise ValueError("runtime is required") + + return ModelInvocationUtils.calculate_tokens( + tenant_id=self.runtime.tenant_id or "", prompt_messages=prompt_messages + ) + + def summary(self, user_id: str, content: str) -> str: + max_tokens = self.get_max_tokens() + + if self.get_prompt_tokens(prompt_messages=[UserPromptMessage(content=content)]) < max_tokens * 0.6: + return content + + def get_prompt_tokens(content: str) -> int: + return self.get_prompt_tokens( + prompt_messages=[SystemPromptMessage(content=_SUMMARY_PROMPT), UserPromptMessage(content=content)] + ) + + def summarize(content: str) -> str: + summary = self.invoke_model( + user_id=user_id, + prompt_messages=[SystemPromptMessage(content=_SUMMARY_PROMPT), UserPromptMessage(content=content)], + stop=[], + ) + + return cast(str, summary.message.content) + + lines = content.split("\n") + new_lines = [] + # split long line into multiple lines + for i in range(len(lines)): + line = lines[i] + if not line.strip(): + continue + if len(line) < max_tokens * 0.5: + new_lines.append(line) + elif get_prompt_tokens(line) > max_tokens * 0.7: + while get_prompt_tokens(line) > max_tokens * 0.7: + new_lines.append(line[: int(max_tokens * 0.5)]) + line = line[int(max_tokens * 0.5) :] + new_lines.append(line) + else: + new_lines.append(line) + + # merge lines into messages with max tokens + messages: list[str] = [] + for j in new_lines: + if len(messages) == 0: + messages.append(j) + else: + if len(messages[-1]) + len(j) < max_tokens * 0.5: + messages[-1] += j + if get_prompt_tokens(messages[-1] + j) > max_tokens * 0.7: + messages.append(j) + else: + messages[-1] += j + + summaries = [] + for i in range(len(messages)): + message = messages[i] + summary = summarize(message) + summaries.append(summary) + + result = "\n".join(summaries) + + if self.get_prompt_tokens(prompt_messages=[UserPromptMessage(content=result)]) > max_tokens * 0.7: + return self.summary(user_id=user_id, content=result) + + return result + + def get_url(self, url: str, user_agent: Optional[str] = None) -> str: + """ + get url + """ + return get_url(url, user_agent=user_agent) diff --git a/api/core/tools/tool/dataset_retriever/dataset_multi_retriever_tool.py b/api/core/tools/tool/dataset_retriever/dataset_multi_retriever_tool.py new file mode 100644 index 0000000000000000000000000000000000000000..a4afea4b9df429d8e907c6e53b1ad79c63edff3d --- /dev/null +++ b/api/core/tools/tool/dataset_retriever/dataset_multi_retriever_tool.py @@ -0,0 +1,196 @@ +import threading +from typing import Any + +from flask import Flask, current_app +from pydantic import BaseModel, Field + +from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler +from core.model_manager import ModelManager +from core.model_runtime.entities.model_entities import ModelType +from core.rag.datasource.retrieval_service import RetrievalService +from core.rag.models.document import Document as RagDocument +from core.rag.rerank.rerank_model import RerankModelRunner +from core.rag.retrieval.retrieval_methods import RetrievalMethod +from core.tools.tool.dataset_retriever.dataset_retriever_base_tool import DatasetRetrieverBaseTool +from extensions.ext_database import db +from models.dataset import Dataset, Document, DocumentSegment + +default_retrieval_model: dict[str, Any] = { + "search_method": RetrievalMethod.SEMANTIC_SEARCH.value, + "reranking_enable": False, + "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""}, + "top_k": 2, + "score_threshold_enabled": False, +} + + +class DatasetMultiRetrieverToolInput(BaseModel): + query: str = Field(..., description="dataset multi retriever and rerank") + + +class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool): + """Tool for querying multi dataset.""" + + name: str = "dataset_" + args_schema: type[BaseModel] = DatasetMultiRetrieverToolInput + description: str = "dataset multi retriever and rerank. " + dataset_ids: list[str] + reranking_provider_name: str + reranking_model_name: str + + @classmethod + def from_dataset(cls, dataset_ids: list[str], tenant_id: str, **kwargs): + return cls( + name=f"dataset_{tenant_id.replace('-', '_')}", tenant_id=tenant_id, dataset_ids=dataset_ids, **kwargs + ) + + def _run(self, query: str) -> str: + threads = [] + all_documents: list[RagDocument] = [] + for dataset_id in self.dataset_ids: + retrieval_thread = threading.Thread( + target=self._retriever, + kwargs={ + "flask_app": current_app._get_current_object(), # type: ignore + "dataset_id": dataset_id, + "query": query, + "all_documents": all_documents, + "hit_callbacks": self.hit_callbacks, + }, + ) + threads.append(retrieval_thread) + retrieval_thread.start() + for thread in threads: + thread.join() + # do rerank for searched documents + model_manager = ModelManager() + rerank_model_instance = model_manager.get_model_instance( + tenant_id=self.tenant_id, + provider=self.reranking_provider_name, + model_type=ModelType.RERANK, + model=self.reranking_model_name, + ) + + rerank_runner = RerankModelRunner(rerank_model_instance) + all_documents = rerank_runner.run(query, all_documents, self.score_threshold, self.top_k) + + for hit_callback in self.hit_callbacks: + hit_callback.on_tool_end(all_documents) + + document_score_list = {} + for item in all_documents: + if item.metadata and item.metadata.get("score"): + document_score_list[item.metadata["doc_id"]] = item.metadata["score"] + + document_context_list = [] + index_node_ids = [document.metadata["doc_id"] for document in all_documents if document.metadata] + segments = DocumentSegment.query.filter( + DocumentSegment.dataset_id.in_(self.dataset_ids), + DocumentSegment.completed_at.isnot(None), + DocumentSegment.status == "completed", + DocumentSegment.enabled == True, + DocumentSegment.index_node_id.in_(index_node_ids), + ).all() + + if segments: + index_node_id_to_position = {id: position for position, id in enumerate(index_node_ids)} + sorted_segments = sorted( + segments, key=lambda segment: index_node_id_to_position.get(segment.index_node_id, float("inf")) + ) + for segment in sorted_segments: + if segment.answer: + document_context_list.append(f"question:{segment.get_sign_content()} answer:{segment.answer}") + else: + document_context_list.append(segment.get_sign_content()) + if self.return_resource: + context_list = [] + resource_number = 1 + for segment in sorted_segments: + dataset = Dataset.query.filter_by(id=segment.dataset_id).first() + document = Document.query.filter( + Document.id == segment.document_id, + Document.enabled == True, + Document.archived == False, + ).first() + if dataset and document: + source = { + "position": resource_number, + "dataset_id": dataset.id, + "dataset_name": dataset.name, + "document_id": document.id, + "document_name": document.name, + "data_source_type": document.data_source_type, + "segment_id": segment.id, + "retriever_from": self.retriever_from, + "score": document_score_list.get(segment.index_node_id, None), + } + + if self.retriever_from == "dev": + source["hit_count"] = segment.hit_count + source["word_count"] = segment.word_count + source["segment_position"] = segment.position + source["index_node_hash"] = segment.index_node_hash + if segment.answer: + source["content"] = f"question:{segment.content} \nanswer:{segment.answer}" + else: + source["content"] = segment.content + context_list.append(source) + resource_number += 1 + + for hit_callback in self.hit_callbacks: + hit_callback.return_retriever_resource_info(context_list) + + return str("\n".join(document_context_list)) + return "" + + def _retriever( + self, + flask_app: Flask, + dataset_id: str, + query: str, + all_documents: list, + hit_callbacks: list[DatasetIndexToolCallbackHandler], + ): + with flask_app.app_context(): + dataset = ( + db.session.query(Dataset).filter(Dataset.tenant_id == self.tenant_id, Dataset.id == dataset_id).first() + ) + + if not dataset: + return [] + + for hit_callback in hit_callbacks: + hit_callback.on_query(query, dataset.id) + + # get retrieval model , if the model is not setting , using default + retrieval_model = dataset.retrieval_model or default_retrieval_model + + if dataset.indexing_technique == "economy": + # use keyword table query + documents = RetrievalService.retrieve( + retrieval_method="keyword_search", + dataset_id=dataset.id, + query=query, + top_k=retrieval_model.get("top_k") or 2, + ) + if documents: + all_documents.extend(documents) + else: + if self.top_k > 0: + # retrieval source + documents = RetrievalService.retrieve( + retrieval_method=retrieval_model["search_method"], + dataset_id=dataset.id, + query=query, + top_k=retrieval_model.get("top_k") or 2, + score_threshold=retrieval_model.get("score_threshold", 0.0) + if retrieval_model["score_threshold_enabled"] + else 0.0, + reranking_model=retrieval_model.get("reranking_model", None) + if retrieval_model["reranking_enable"] + else None, + reranking_mode=retrieval_model.get("reranking_mode") or "reranking_model", + weights=retrieval_model.get("weights", None), + ) + + all_documents.extend(documents) diff --git a/api/core/tools/tool/dataset_retriever/dataset_retriever_base_tool.py b/api/core/tools/tool/dataset_retriever/dataset_retriever_base_tool.py new file mode 100644 index 0000000000000000000000000000000000000000..a4d2de3b1c8ef3fbbde6476e4ea430fa0ce082f7 --- /dev/null +++ b/api/core/tools/tool/dataset_retriever/dataset_retriever_base_tool.py @@ -0,0 +1,33 @@ +from abc import abstractmethod +from typing import Any, Optional + +from msal_extensions.persistence import ABC # type: ignore +from pydantic import BaseModel, ConfigDict + +from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler + + +class DatasetRetrieverBaseTool(BaseModel, ABC): + """Tool for querying a Dataset.""" + + name: str = "dataset" + description: str = "use this to retrieve a dataset. " + tenant_id: str + top_k: int = 2 + score_threshold: Optional[float] = None + hit_callbacks: list[DatasetIndexToolCallbackHandler] = [] + return_resource: bool + retriever_from: str + model_config = ConfigDict(arbitrary_types_allowed=True) + + @abstractmethod + def _run( + self, + *args: Any, + **kwargs: Any, + ) -> Any: + """Use the tool. + + Add run_manager: Optional[CallbackManagerForToolRun] = None + to child implementations to enable tracing, + """ diff --git a/api/core/tools/tool/dataset_retriever/dataset_retriever_tool.py b/api/core/tools/tool/dataset_retriever/dataset_retriever_tool.py new file mode 100644 index 0000000000000000000000000000000000000000..b382016473055d2a4ee13dd66935ddc9094c4edf --- /dev/null +++ b/api/core/tools/tool/dataset_retriever/dataset_retriever_tool.py @@ -0,0 +1,196 @@ +from typing import Any + +from pydantic import BaseModel, Field + +from core.rag.datasource.retrieval_service import RetrievalService +from core.rag.models.document import Document as RetrievalDocument +from core.rag.retrieval.retrieval_methods import RetrievalMethod +from core.tools.tool.dataset_retriever.dataset_retriever_base_tool import DatasetRetrieverBaseTool +from extensions.ext_database import db +from models.dataset import Dataset, Document, DocumentSegment +from services.external_knowledge_service import ExternalDatasetService + +default_retrieval_model = { + "search_method": RetrievalMethod.SEMANTIC_SEARCH.value, + "reranking_enable": False, + "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""}, + "reranking_mode": "reranking_model", + "top_k": 2, + "score_threshold_enabled": False, +} + + +class DatasetRetrieverToolInput(BaseModel): + query: str = Field(..., description="Query for the dataset to be used to retrieve the dataset.") + + +class DatasetRetrieverTool(DatasetRetrieverBaseTool): + """Tool for querying a Dataset.""" + + name: str = "dataset" + args_schema: type[BaseModel] = DatasetRetrieverToolInput + description: str = "use this to retrieve a dataset. " + dataset_id: str + + @classmethod + def from_dataset(cls, dataset: Dataset, **kwargs): + description = dataset.description + if not description: + description = "useful for when you want to answer queries about the " + dataset.name + + description = description.replace("\n", "").replace("\r", "") + return cls( + name=f"dataset_{dataset.id.replace('-', '_')}", + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + description=description, + **kwargs, + ) + + def _run(self, query: str) -> str: + dataset = ( + db.session.query(Dataset).filter(Dataset.tenant_id == self.tenant_id, Dataset.id == self.dataset_id).first() + ) + + if not dataset: + return "" + + for hit_callback in self.hit_callbacks: + hit_callback.on_query(query, dataset.id) + if dataset.provider == "external": + results = [] + external_documents = ExternalDatasetService.fetch_external_knowledge_retrieval( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + query=query, + external_retrieval_parameters=dataset.retrieval_model, + ) + for external_document in external_documents: + document = RetrievalDocument( + page_content=external_document.get("content"), + metadata=external_document.get("metadata"), + provider="external", + ) + if document.metadata is not None: + document.metadata["score"] = external_document.get("score") + document.metadata["title"] = external_document.get("title") + document.metadata["dataset_id"] = dataset.id + document.metadata["dataset_name"] = dataset.name + results.append(document) + # deal with external documents + context_list = [] + for position, item in enumerate(results, start=1): + if item.metadata is not None: + source = { + "position": position, + "dataset_id": item.metadata.get("dataset_id"), + "dataset_name": item.metadata.get("dataset_name"), + "document_name": item.metadata.get("title"), + "data_source_type": "external", + "retriever_from": self.retriever_from, + "score": item.metadata.get("score"), + "title": item.metadata.get("title"), + "content": item.page_content, + } + context_list.append(source) + for hit_callback in self.hit_callbacks: + hit_callback.return_retriever_resource_info(context_list) + + return str("\n".join([item.page_content for item in results])) + else: + # get retrieval model , if the model is not setting , using default + retrieval_model: dict[str, Any] = dataset.retrieval_model or default_retrieval_model + if dataset.indexing_technique == "economy": + # use keyword table query + documents = RetrievalService.retrieve( + retrieval_method="keyword_search", dataset_id=dataset.id, query=query, top_k=self.top_k + ) + return str("\n".join([document.page_content for document in documents])) + else: + if self.top_k > 0: + # retrieval source + documents = RetrievalService.retrieve( + retrieval_method=retrieval_model.get("search_method", "semantic_search"), + dataset_id=dataset.id, + query=query, + top_k=self.top_k, + score_threshold=retrieval_model.get("score_threshold", 0.0) + if retrieval_model["score_threshold_enabled"] + else 0.0, + reranking_model=retrieval_model.get("reranking_model") + if retrieval_model["reranking_enable"] + else None, + reranking_mode=retrieval_model.get("reranking_mode") or "reranking_model", + weights=retrieval_model.get("weights"), + ) + else: + documents = [] + + for hit_callback in self.hit_callbacks: + hit_callback.on_tool_end(documents) + document_score_list = {} + if dataset.indexing_technique != "economy": + for item in documents: + if item.metadata is not None and item.metadata.get("score"): + document_score_list[item.metadata["doc_id"]] = item.metadata["score"] + document_context_list = [] + index_node_ids = [document.metadata["doc_id"] for document in documents] + segments = DocumentSegment.query.filter( + DocumentSegment.dataset_id == self.dataset_id, + DocumentSegment.completed_at.isnot(None), + DocumentSegment.status == "completed", + DocumentSegment.enabled == True, + DocumentSegment.index_node_id.in_(index_node_ids), + ).all() + + if segments: + index_node_id_to_position = {id: position for position, id in enumerate(index_node_ids)} + sorted_segments = sorted( + segments, key=lambda segment: index_node_id_to_position.get(segment.index_node_id, float("inf")) + ) + for segment in sorted_segments: + if segment.answer: + document_context_list.append( + f"question:{segment.get_sign_content()} answer:{segment.answer}" + ) + else: + document_context_list.append(segment.get_sign_content()) + if self.return_resource: + context_list = [] + resource_number = 1 + for segment in sorted_segments: + document_segment = Document.query.filter( + Document.id == segment.document_id, + Document.enabled == True, + Document.archived == False, + ).first() + if not document_segment: + continue + if dataset and document_segment: + source = { + "position": resource_number, + "dataset_id": dataset.id, + "dataset_name": dataset.name, + "document_id": document_segment.id, + "document_name": document_segment.name, + "data_source_type": document_segment.data_source_type, + "segment_id": segment.id, + "retriever_from": self.retriever_from, + "score": document_score_list.get(segment.index_node_id, None), + } + if self.retriever_from == "dev": + source["hit_count"] = segment.hit_count + source["word_count"] = segment.word_count + source["segment_position"] = segment.position + source["index_node_hash"] = segment.index_node_hash + if segment.answer: + source["content"] = f"question:{segment.content} \nanswer:{segment.answer}" + else: + source["content"] = segment.content + context_list.append(source) + resource_number += 1 + + for hit_callback in self.hit_callbacks: + hit_callback.return_retriever_resource_info(context_list) + + return str("\n".join(document_context_list)) diff --git a/api/core/tools/tool/dataset_retriever_tool.py b/api/core/tools/tool/dataset_retriever_tool.py new file mode 100644 index 0000000000000000000000000000000000000000..2d7e193e152645ca70ec41051143376d76642d04 --- /dev/null +++ b/api/core/tools/tool/dataset_retriever_tool.py @@ -0,0 +1,114 @@ +from typing import Any, Optional + +from core.app.app_config.entities import DatasetRetrieveConfigEntity +from core.app.entities.app_invoke_entities import InvokeFrom +from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler +from core.rag.retrieval.dataset_retrieval import DatasetRetrieval +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_entities import ( + ToolDescription, + ToolIdentity, + ToolInvokeMessage, + ToolParameter, + ToolProviderType, +) +from core.tools.tool.dataset_retriever.dataset_retriever_base_tool import DatasetRetrieverBaseTool +from core.tools.tool.tool import Tool + + +class DatasetRetrieverTool(Tool): + retrieval_tool: DatasetRetrieverBaseTool + + @staticmethod + def get_dataset_tools( + tenant_id: str, + dataset_ids: list[str], + retrieve_config: Optional[DatasetRetrieveConfigEntity], + return_resource: bool, + invoke_from: InvokeFrom, + hit_callback: DatasetIndexToolCallbackHandler, + ) -> list["DatasetRetrieverTool"]: + """ + get dataset tool + """ + # check if retrieve_config is valid + if dataset_ids is None or len(dataset_ids) == 0: + return [] + if retrieve_config is None: + return [] + + feature = DatasetRetrieval() + + # save original retrieve strategy, and set retrieve strategy to SINGLE + # Agent only support SINGLE mode + original_retriever_mode = retrieve_config.retrieve_strategy + retrieve_config.retrieve_strategy = DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE + retrieval_tools = feature.to_dataset_retriever_tool( + tenant_id=tenant_id, + dataset_ids=dataset_ids, + retrieve_config=retrieve_config, + return_resource=return_resource, + invoke_from=invoke_from, + hit_callback=hit_callback, + ) + if retrieval_tools is None: + return [] + # restore retrieve strategy + retrieve_config.retrieve_strategy = original_retriever_mode + + # convert retrieval tools to Tools + tools = [] + for retrieval_tool in retrieval_tools: + tool = DatasetRetrieverTool( + retrieval_tool=retrieval_tool, + identity=ToolIdentity( + provider="", author="", name=retrieval_tool.name, label=I18nObject(en_US="", zh_Hans="") + ), + parameters=[], + is_team_authorization=True, + description=ToolDescription(human=I18nObject(en_US="", zh_Hans=""), llm=retrieval_tool.description), + runtime=DatasetRetrieverTool.Runtime(), + ) + + tools.append(tool) + + return tools + + def get_runtime_parameters(self) -> list[ToolParameter]: + return [ + ToolParameter( + name="query", + label=I18nObject(en_US="", zh_Hans=""), + human_description=I18nObject(en_US="", zh_Hans=""), + type=ToolParameter.ToolParameterType.STRING, + form=ToolParameter.ToolParameterForm.LLM, + llm_description="Query for the dataset to be used to retrieve the dataset.", + required=True, + default="", + placeholder=I18nObject(en_US="", zh_Hans=""), + ), + ] + + def tool_provider_type(self) -> ToolProviderType: + return ToolProviderType.DATASET_RETRIEVAL + + def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage | list[ToolInvokeMessage]: + """ + invoke dataset retriever tool + """ + query = tool_parameters.get("query") + if not query: + return self.create_text_message(text="please input query") + + # invoke dataset retriever tool + result = self.retrieval_tool._run(query=query) + + return self.create_text_message(text=result) + + def validate_credentials( + self, credentials: dict[str, Any], parameters: dict[str, Any], format_only: bool = False + ) -> str | None: + """ + validate the credentials for dataset retriever tool + """ + pass diff --git a/api/core/tools/tool/tool.py b/api/core/tools/tool/tool.py new file mode 100644 index 0000000000000000000000000000000000000000..4094207beb5c9f43f3346a41089969fc9b900eb2 --- /dev/null +++ b/api/core/tools/tool/tool.py @@ -0,0 +1,355 @@ +from abc import ABC, abstractmethod +from collections.abc import Mapping +from copy import deepcopy +from enum import Enum, StrEnum +from typing import TYPE_CHECKING, Any, Optional, Union + +from pydantic import BaseModel, ConfigDict, field_validator +from pydantic_core.core_schema import ValidationInfo + +from core.app.entities.app_invoke_entities import InvokeFrom +from core.tools.entities.tool_entities import ( + ToolDescription, + ToolIdentity, + ToolInvokeFrom, + ToolInvokeMessage, + ToolParameter, + ToolProviderType, + ToolRuntimeImageVariable, + ToolRuntimeVariable, + ToolRuntimeVariablePool, +) +from core.tools.tool_file_manager import ToolFileManager + +if TYPE_CHECKING: + from core.file.models import File + + +class Tool(BaseModel, ABC): + identity: Optional[ToolIdentity] = None + parameters: Optional[list[ToolParameter]] = None + description: Optional[ToolDescription] = None + is_team_authorization: bool = False + + # pydantic configs + model_config = ConfigDict(protected_namespaces=()) + + @field_validator("parameters", mode="before") + @classmethod + def set_parameters(cls, v, validation_info: ValidationInfo) -> list[ToolParameter]: + return v or [] + + class Runtime(BaseModel): + """ + Meta data of a tool call processing + """ + + def __init__(self, **data: Any): + super().__init__(**data) + if not self.runtime_parameters: + self.runtime_parameters = {} + + tenant_id: Optional[str] = None + tool_id: Optional[str] = None + invoke_from: Optional[InvokeFrom] = None + tool_invoke_from: Optional[ToolInvokeFrom] = None + credentials: Optional[dict[str, Any]] = None + runtime_parameters: Optional[dict[str, Any]] = None + + runtime: Optional[Runtime] = None + variables: Optional[ToolRuntimeVariablePool] = None + + def __init__(self, **data: Any): + super().__init__(**data) + + class VariableKey(StrEnum): + IMAGE = "image" + DOCUMENT = "document" + VIDEO = "video" + AUDIO = "audio" + CUSTOM = "custom" + + def fork_tool_runtime(self, runtime: dict[str, Any]) -> "Tool": + """ + fork a new tool with meta data + + :param meta: the meta data of a tool call processing, tenant_id is required + :return: the new tool + """ + return self.__class__( + identity=self.identity.model_copy() if self.identity else None, + parameters=self.parameters.copy() if self.parameters else None, + description=self.description.model_copy() if self.description else None, + runtime=Tool.Runtime(**runtime), + ) + + @abstractmethod + def tool_provider_type(self) -> ToolProviderType: + """ + get the tool provider type + + :return: the tool provider type + """ + + def load_variables(self, variables: ToolRuntimeVariablePool | None) -> None: + """ + load variables from database + + :param conversation_id: the conversation id + """ + self.variables = variables + + def set_image_variable(self, variable_name: str, image_key: str) -> None: + """ + set an image variable + """ + if not self.variables: + return + if self.identity is None: + return + + self.variables.set_file(self.identity.name, variable_name, image_key) + + def set_text_variable(self, variable_name: str, text: str) -> None: + """ + set a text variable + """ + if not self.variables: + return + if self.identity is None: + return + + self.variables.set_text(self.identity.name, variable_name, text) + + def get_variable(self, name: Union[str, Enum]) -> Optional[ToolRuntimeVariable]: + """ + get a variable + + :param name: the name of the variable + :return: the variable + """ + if not self.variables: + return None + + if isinstance(name, Enum): + name = name.value + + for variable in self.variables.pool: + if variable.name == name: + return variable + + return None + + def get_default_image_variable(self) -> Optional[ToolRuntimeVariable]: + """ + get the default image variable + + :return: the image variable + """ + if not self.variables: + return None + + return self.get_variable(self.VariableKey.IMAGE) + + def get_variable_file(self, name: Union[str, Enum]) -> Optional[bytes]: + """ + get a variable file + + :param name: the name of the variable + :return: the variable file + """ + variable = self.get_variable(name) + if not variable: + return None + + if not isinstance(variable, ToolRuntimeImageVariable): + return None + + message_file_id = variable.value + # get file binary + file_binary = ToolFileManager.get_file_binary_by_message_file_id(message_file_id) + if not file_binary: + return None + + return file_binary[0] + + def list_variables(self) -> list[ToolRuntimeVariable]: + """ + list all variables + + :return: the variables + """ + if not self.variables: + return [] + + return self.variables.pool + + def list_default_image_variables(self) -> list[ToolRuntimeVariable]: + """ + list all image variables + + :return: the image variables + """ + if not self.variables: + return [] + + result = [] + + for variable in self.variables.pool: + if variable.name.startswith(self.VariableKey.IMAGE.value): + result.append(variable) + + return result + + def invoke(self, user_id: str, tool_parameters: Mapping[str, Any]) -> list[ToolInvokeMessage]: + # update tool_parameters + # TODO: Fix type error. + if self.runtime is None: + return [] + if self.runtime.runtime_parameters: + # Convert Mapping to dict before updating + tool_parameters = dict(tool_parameters) + tool_parameters.update(self.runtime.runtime_parameters) + + # try parse tool parameters into the correct type + tool_parameters = self._transform_tool_parameters_type(tool_parameters) + + result = self._invoke( + user_id=user_id, + tool_parameters=tool_parameters, + ) + + if not isinstance(result, list): + result = [result] + + if not all(isinstance(message, ToolInvokeMessage) for message in result): + raise ValueError( + f"Invalid return type from {self.__class__.__name__}._invoke method. " + "Expected ToolInvokeMessage or list of ToolInvokeMessage." + ) + + return result + + def _transform_tool_parameters_type(self, tool_parameters: Mapping[str, Any]) -> dict[str, Any]: + """ + Transform tool parameters type + """ + # Temp fix for the issue that the tool parameters will be converted to empty while validating the credentials + result: dict[str, Any] = deepcopy(dict(tool_parameters)) + for parameter in self.parameters or []: + if parameter.name in tool_parameters: + result[parameter.name] = parameter.type.cast_value(tool_parameters[parameter.name]) + + return result + + @abstractmethod + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + pass + + def validate_credentials( + self, credentials: dict[str, Any], parameters: dict[str, Any], format_only: bool = False + ) -> str | None: + """ + validate the credentials + + :param credentials: the credentials + :param parameters: the parameters + :param format_only: only return the formatted + """ + pass + + def get_runtime_parameters(self) -> list[ToolParameter]: + """ + get the runtime parameters + + interface for developer to dynamic change the parameters of a tool depends on the variables pool + + :return: the runtime parameters + """ + return self.parameters or [] + + def get_all_runtime_parameters(self) -> list[ToolParameter]: + """ + get all runtime parameters + + :return: all runtime parameters + """ + parameters = self.parameters or [] + parameters = parameters.copy() + user_parameters = self.get_runtime_parameters() + user_parameters = user_parameters.copy() + + # override parameters + for parameter in user_parameters: + # check if parameter in tool parameters + found = False + for tool_parameter in parameters: + if tool_parameter.name == parameter.name: + found = True + break + + if found: + # override parameter + tool_parameter.type = parameter.type + tool_parameter.form = parameter.form + tool_parameter.required = parameter.required + tool_parameter.default = parameter.default + tool_parameter.options = parameter.options + tool_parameter.llm_description = parameter.llm_description + else: + # add new parameter + parameters.append(parameter) + + return parameters + + def create_image_message(self, image: str, save_as: str = "") -> ToolInvokeMessage: + """ + create an image message + + :param image: the url of the image + :return: the image message + """ + return ToolInvokeMessage(type=ToolInvokeMessage.MessageType.IMAGE, message=image, save_as=save_as) + + def create_file_message(self, file: "File") -> ToolInvokeMessage: + return ToolInvokeMessage(type=ToolInvokeMessage.MessageType.FILE, message="", meta={"file": file}, save_as="") + + def create_link_message(self, link: str, save_as: str = "") -> ToolInvokeMessage: + """ + create a link message + + :param link: the url of the link + :return: the link message + """ + return ToolInvokeMessage(type=ToolInvokeMessage.MessageType.LINK, message=link, save_as=save_as) + + def create_text_message(self, text: str, save_as: str = "") -> ToolInvokeMessage: + """ + create a text message + + :param text: the text + :return: the text message + """ + return ToolInvokeMessage(type=ToolInvokeMessage.MessageType.TEXT, message=text, save_as=save_as) + + def create_blob_message(self, blob: bytes, meta: Optional[dict] = None, save_as: str = "") -> ToolInvokeMessage: + """ + create a blob message + + :param blob: the blob + :return: the blob message + """ + return ToolInvokeMessage( + type=ToolInvokeMessage.MessageType.BLOB, + message=blob, + meta=meta or {}, + save_as=save_as, + ) + + def create_json_message(self, object: dict) -> ToolInvokeMessage: + """ + create a json message + """ + return ToolInvokeMessage(type=ToolInvokeMessage.MessageType.JSON, message=object) diff --git a/api/core/tools/tool/workflow_tool.py b/api/core/tools/tool/workflow_tool.py new file mode 100644 index 0000000000000000000000000000000000000000..9d13e2c49a55e32542336b5c1c6b9c28c223ad1e --- /dev/null +++ b/api/core/tools/tool/workflow_tool.py @@ -0,0 +1,221 @@ +import json +import logging +from copy import deepcopy +from typing import Any, Optional, Union, cast + +from core.file import FILE_MODEL_IDENTITY, File, FileTransferMethod +from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter, ToolProviderType +from core.tools.tool.tool import Tool +from extensions.ext_database import db +from factories.file_factory import build_from_mapping +from models.account import Account +from models.model import App, EndUser +from models.workflow import Workflow + +logger = logging.getLogger(__name__) + + +class WorkflowTool(Tool): + workflow_app_id: str + version: str + workflow_entities: dict[str, Any] + workflow_call_depth: int + thread_pool_id: Optional[str] = None + + label: str + + """ + Workflow tool. + """ + + def tool_provider_type(self) -> ToolProviderType: + """ + get the tool provider type + + :return: the tool provider type + """ + return ToolProviderType.WORKFLOW + + def _invoke( + self, user_id: str, tool_parameters: dict[str, Any] + ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: + """ + invoke the tool + """ + app = self._get_app(app_id=self.workflow_app_id) + workflow = self._get_workflow(app_id=self.workflow_app_id, version=self.version) + + # transform the tool parameters + tool_parameters, files = self._transform_args(tool_parameters=tool_parameters) + + from core.app.apps.workflow.app_generator import WorkflowAppGenerator + + generator = WorkflowAppGenerator() + assert self.runtime is not None + assert self.runtime.invoke_from is not None + result = generator.generate( + app_model=app, + workflow=workflow, + user=self._get_user(user_id), + args={"inputs": tool_parameters, "files": files}, + invoke_from=self.runtime.invoke_from, + streaming=False, + call_depth=self.workflow_call_depth + 1, + workflow_thread_pool_id=self.thread_pool_id, + ) + assert isinstance(result, dict) + data = result.get("data", {}) + + if data.get("error"): + raise Exception(data.get("error")) + + r = [] + + outputs = data.get("outputs") + if outputs == None: + outputs = {} + else: + outputs, extracted_files = self._extract_files(outputs) + for f in extracted_files: + r.append(self.create_file_message(f)) + + r.append(self.create_text_message(json.dumps(outputs, ensure_ascii=False))) + r.append(self.create_json_message(outputs)) + + return r + + def _get_user(self, user_id: str) -> Union[EndUser, Account]: + """ + get the user by user id + """ + + user = db.session.query(EndUser).filter(EndUser.id == user_id).first() + if not user: + user = db.session.query(Account).filter(Account.id == user_id).first() + + if not user: + raise ValueError("user not found") + + return user + + def fork_tool_runtime(self, runtime: dict[str, Any]) -> "WorkflowTool": + """ + fork a new tool with meta data + + :param meta: the meta data of a tool call processing, tenant_id is required + :return: the new tool + """ + return self.__class__( + identity=deepcopy(self.identity), + parameters=deepcopy(self.parameters), + description=deepcopy(self.description), + runtime=Tool.Runtime(**runtime), + workflow_app_id=self.workflow_app_id, + workflow_entities=self.workflow_entities, + workflow_call_depth=self.workflow_call_depth, + version=self.version, + label=self.label, + ) + + def _get_workflow(self, app_id: str, version: str) -> Workflow: + """ + get the workflow by app id and version + """ + if not version: + workflow = ( + db.session.query(Workflow) + .filter(Workflow.app_id == app_id, Workflow.version != "draft") + .order_by(Workflow.created_at.desc()) + .first() + ) + else: + workflow = db.session.query(Workflow).filter(Workflow.app_id == app_id, Workflow.version == version).first() + + if not workflow: + raise ValueError("workflow not found or not published") + + return workflow + + def _get_app(self, app_id: str) -> App: + """ + get the app by app id + """ + app = db.session.query(App).filter(App.id == app_id).first() + if not app: + raise ValueError("app not found") + + return app + + def _transform_args(self, tool_parameters: dict) -> tuple[dict, list[dict]]: + """ + transform the tool parameters + + :param tool_parameters: the tool parameters + :return: tool_parameters, files + """ + parameter_rules = self.get_all_runtime_parameters() + parameters_result = {} + files = [] + for parameter in parameter_rules: + if parameter.type == ToolParameter.ToolParameterType.SYSTEM_FILES: + file = tool_parameters.get(parameter.name) + if file: + try: + file_var_list = [File.model_validate(f) for f in file] + for file in file_var_list: + file_dict: dict[str, str | None] = { + "transfer_method": file.transfer_method.value, + "type": file.type.value, + } + if file.transfer_method == FileTransferMethod.TOOL_FILE: + file_dict["tool_file_id"] = file.related_id + elif file.transfer_method == FileTransferMethod.LOCAL_FILE: + file_dict["upload_file_id"] = file.related_id + elif file.transfer_method == FileTransferMethod.REMOTE_URL: + file_dict["url"] = file.generate_url() + + files.append(file_dict) + except Exception as e: + logger.exception(f"Failed to transform file {file}") + else: + parameters_result[parameter.name] = tool_parameters.get(parameter.name) + + return parameters_result, files + + def _extract_files(self, outputs: dict) -> tuple[dict, list[File]]: + """ + extract files from the result + + :param result: the result + :return: the result, files + """ + files = [] + result = {} + for key, value in outputs.items(): + if isinstance(value, list): + for item in value: + if isinstance(item, dict) and item.get("dify_model_identity") == FILE_MODEL_IDENTITY: + item = self._update_file_mapping(item) + file = build_from_mapping( + mapping=item, + tenant_id=str(cast(Tool.Runtime, self.runtime).tenant_id), + ) + files.append(file) + elif isinstance(value, dict) and value.get("dify_model_identity") == FILE_MODEL_IDENTITY: + value = self._update_file_mapping(value) + file = build_from_mapping( + mapping=value, + tenant_id=str(cast(Tool.Runtime, self.runtime).tenant_id), + ) + files.append(file) + + result[key] = value + return result, files + + def _update_file_mapping(self, file_dict: dict) -> dict: + transfer_method = FileTransferMethod.value_of(file_dict.get("transfer_method")) + if transfer_method == FileTransferMethod.TOOL_FILE: + file_dict["tool_file_id"] = file_dict.get("related_id") + elif transfer_method == FileTransferMethod.LOCAL_FILE: + file_dict["upload_file_id"] = file_dict.get("related_id") + return file_dict diff --git a/api/core/tools/tool_engine.py b/api/core/tools/tool_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..f7a8ed63f401d53f5cd5dc79405654bb65837638 --- /dev/null +++ b/api/core/tools/tool_engine.py @@ -0,0 +1,323 @@ +import json +from collections.abc import Mapping +from copy import deepcopy +from datetime import UTC, datetime +from mimetypes import guess_type +from typing import Any, Optional, Union, cast + +from yarl import URL + +from core.app.entities.app_invoke_entities import InvokeFrom +from core.callback_handler.agent_tool_callback_handler import DifyAgentCallbackHandler +from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler +from core.file import FileType +from core.file.models import FileTransferMethod +from core.ops.ops_trace_manager import TraceQueueManager +from core.tools.entities.tool_entities import ToolInvokeMessage, ToolInvokeMessageBinary, ToolInvokeMeta, ToolParameter +from core.tools.errors import ( + ToolEngineInvokeError, + ToolInvokeError, + ToolNotFoundError, + ToolNotSupportedError, + ToolParameterValidationError, + ToolProviderCredentialValidationError, + ToolProviderNotFoundError, +) +from core.tools.tool.tool import Tool +from core.tools.tool.workflow_tool import WorkflowTool +from core.tools.utils.message_transformer import ToolFileMessageTransformer +from extensions.ext_database import db +from models.enums import CreatedByRole +from models.model import Message, MessageFile + + +class ToolEngine: + """ + Tool runtime engine take care of the tool executions. + """ + + @staticmethod + def agent_invoke( + tool: Tool, + tool_parameters: Union[str, dict], + user_id: str, + tenant_id: str, + message: Message, + invoke_from: InvokeFrom, + agent_tool_callback: DifyAgentCallbackHandler, + trace_manager: Optional[TraceQueueManager] = None, + ) -> tuple[str, list[tuple[MessageFile, str]], ToolInvokeMeta]: + """ + Agent invokes the tool with the given arguments. + """ + # check if arguments is a string + if isinstance(tool_parameters, str): + # check if this tool has only one parameter + parameters = [ + parameter + for parameter in tool.get_runtime_parameters() + if parameter.form == ToolParameter.ToolParameterForm.LLM + ] + if parameters and len(parameters) == 1: + tool_parameters = {parameters[0].name: tool_parameters} + else: + try: + tool_parameters = json.loads(tool_parameters) + except Exception as e: + pass + if not isinstance(tool_parameters, dict): + raise ValueError(f"tool_parameters should be a dict, but got a string: {tool_parameters}") + + # invoke the tool + if tool.identity is None: + raise ValueError("tool identity is not set") + try: + # hit the callback handler + agent_tool_callback.on_tool_start(tool_name=tool.identity.name, tool_inputs=tool_parameters) + + meta, response = ToolEngine._invoke(tool, tool_parameters, user_id) + response = ToolFileMessageTransformer.transform_tool_invoke_messages( + messages=response, user_id=user_id, tenant_id=tenant_id, conversation_id=message.conversation_id + ) + + # extract binary data from tool invoke message + binary_files = ToolEngine._extract_tool_response_binary(response) + # create message file + message_files = ToolEngine._create_message_files( + tool_messages=binary_files, agent_message=message, invoke_from=invoke_from, user_id=user_id + ) + + plain_text = ToolEngine._convert_tool_response_to_str(response) + + # hit the callback handler + agent_tool_callback.on_tool_end( + tool_name=tool.identity.name, + tool_inputs=tool_parameters, + tool_outputs=plain_text, + message_id=message.id, + trace_manager=trace_manager, + ) + + # transform tool invoke message to get LLM friendly message + return plain_text, message_files, meta + except ToolProviderCredentialValidationError as e: + error_response = "Please check your tool provider credentials" + agent_tool_callback.on_tool_error(e) + except (ToolNotFoundError, ToolNotSupportedError, ToolProviderNotFoundError) as e: + error_response = f"there is not a tool named {tool.identity.name}" + agent_tool_callback.on_tool_error(e) + except ToolParameterValidationError as e: + error_response = f"tool parameters validation error: {e}, please check your tool parameters" + agent_tool_callback.on_tool_error(e) + except ToolInvokeError as e: + error_response = f"tool invoke error: {e}" + agent_tool_callback.on_tool_error(e) + except ToolEngineInvokeError as e: + meta = e.meta + error_response = f"tool invoke error: {meta.error}" + agent_tool_callback.on_tool_error(e) + return error_response, [], meta + except Exception as e: + error_response = f"unknown error: {e}" + agent_tool_callback.on_tool_error(e) + + return error_response, [], ToolInvokeMeta.error_instance(error_response) + + @staticmethod + def workflow_invoke( + tool: Tool, + tool_parameters: Mapping[str, Any], + user_id: str, + workflow_tool_callback: DifyWorkflowCallbackHandler, + workflow_call_depth: int, + thread_pool_id: Optional[str] = None, + ) -> list[ToolInvokeMessage]: + """ + Workflow invokes the tool with the given arguments. + """ + try: + # hit the callback handler + assert tool.identity is not None + workflow_tool_callback.on_tool_start(tool_name=tool.identity.name, tool_inputs=tool_parameters) + + if isinstance(tool, WorkflowTool): + tool.workflow_call_depth = workflow_call_depth + 1 + tool.thread_pool_id = thread_pool_id + + if tool.runtime and tool.runtime.runtime_parameters: + tool_parameters = {**tool.runtime.runtime_parameters, **tool_parameters} + response = tool.invoke(user_id=user_id, tool_parameters=tool_parameters) + + # hit the callback handler + workflow_tool_callback.on_tool_end( + tool_name=tool.identity.name, + tool_inputs=tool_parameters, + tool_outputs=response, + ) + + return response + except Exception as e: + workflow_tool_callback.on_tool_error(e) + raise e + + @staticmethod + def _invoke(tool: Tool, tool_parameters: dict, user_id: str) -> tuple[ToolInvokeMeta, list[ToolInvokeMessage]]: + """ + Invoke the tool with the given arguments. + """ + if tool.identity is None: + raise ValueError("tool identity is not set") + started_at = datetime.now(UTC) + meta = ToolInvokeMeta( + time_cost=0.0, + error=None, + tool_config={ + "tool_name": tool.identity.name, + "tool_provider": tool.identity.provider, + "tool_provider_type": tool.tool_provider_type().value, + "tool_parameters": deepcopy(tool.runtime.runtime_parameters) if tool.runtime else {}, + "tool_icon": tool.identity.icon, + }, + ) + try: + response = tool.invoke(user_id, tool_parameters) + except Exception as e: + meta.error = str(e) + raise ToolEngineInvokeError(meta) + finally: + ended_at = datetime.now(UTC) + meta.time_cost = (ended_at - started_at).total_seconds() + + return meta, response + + @staticmethod + def _convert_tool_response_to_str(tool_response: list[ToolInvokeMessage]) -> str: + """ + Handle tool response + """ + result = "" + for response in tool_response: + if response.type == ToolInvokeMessage.MessageType.TEXT: + result += str(response.message) if response.message is not None else "" + elif response.type == ToolInvokeMessage.MessageType.LINK: + result += f"result link: {response.message!r}. please tell user to check it." + elif response.type in {ToolInvokeMessage.MessageType.IMAGE_LINK, ToolInvokeMessage.MessageType.IMAGE}: + result += ( + "image has been created and sent to user already, you do not need to create it," + " just tell the user to check it now." + ) + elif response.type == ToolInvokeMessage.MessageType.JSON: + result += f"tool response: {json.dumps(response.message, ensure_ascii=False)}." + else: + result += f"tool response: {response.message!r}." + + return result + + @staticmethod + def _extract_tool_response_binary(tool_response: list[ToolInvokeMessage]) -> list[ToolInvokeMessageBinary]: + """ + Extract tool response binary + """ + result = [] + + for response in tool_response: + if response.type in {ToolInvokeMessage.MessageType.IMAGE_LINK, ToolInvokeMessage.MessageType.IMAGE}: + mimetype = None + if response.meta.get("mime_type"): + mimetype = response.meta.get("mime_type") + else: + try: + url = URL(cast(str, response.message)) + extension = url.suffix + guess_type_result, _ = guess_type(f"a{extension}") + if guess_type_result: + mimetype = guess_type_result + except Exception: + pass + + if not mimetype: + mimetype = "image/jpeg" + + result.append( + ToolInvokeMessageBinary( + mimetype=response.meta.get("mime_type", "image/jpeg"), + url=cast(str, response.message), + save_as=response.save_as, + ) + ) + elif response.type == ToolInvokeMessage.MessageType.BLOB: + result.append( + ToolInvokeMessageBinary( + mimetype=response.meta.get("mime_type", "octet/stream"), + url=cast(str, response.message), + save_as=response.save_as, + ) + ) + elif response.type == ToolInvokeMessage.MessageType.LINK: + # check if there is a mime type in meta + if response.meta and "mime_type" in response.meta: + result.append( + ToolInvokeMessageBinary( + mimetype=response.meta.get("mime_type", "octet/stream") + if response.meta + else "octet/stream", + url=cast(str, response.message), + save_as=response.save_as, + ) + ) + + return result + + @staticmethod + def _create_message_files( + tool_messages: list[ToolInvokeMessageBinary], + agent_message: Message, + invoke_from: InvokeFrom, + user_id: str, + ) -> list[tuple[Any, str]]: + """ + Create message file + + :param messages: messages + :return: message files, should save as variable + """ + result = [] + + for message in tool_messages: + if "image" in message.mimetype: + file_type = FileType.IMAGE + elif "video" in message.mimetype: + file_type = FileType.VIDEO + elif "audio" in message.mimetype: + file_type = FileType.AUDIO + elif "text" in message.mimetype or "pdf" in message.mimetype: + file_type = FileType.DOCUMENT + else: + file_type = FileType.CUSTOM + + # extract tool file id from url + tool_file_id = message.url.split("/")[-1].split(".")[0] + message_file = MessageFile( + message_id=agent_message.id, + type=file_type, + transfer_method=FileTransferMethod.TOOL_FILE, + belongs_to="assistant", + url=message.url, + upload_file_id=tool_file_id, + created_by_role=( + CreatedByRole.ACCOUNT + if invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER} + else CreatedByRole.END_USER + ), + created_by=user_id, + ) + + db.session.add(message_file) + db.session.commit() + db.session.refresh(message_file) + + result.append((message_file.id, message.save_as)) + + db.session.close() + + return result diff --git a/api/core/tools/tool_file_manager.py b/api/core/tools/tool_file_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..2aaca6d82e36b1ab3fd9d7e3f05688ad8d1c161f --- /dev/null +++ b/api/core/tools/tool_file_manager.py @@ -0,0 +1,223 @@ +import base64 +import hashlib +import hmac +import logging +import os +import time +from mimetypes import guess_extension, guess_type +from typing import Optional, Union +from uuid import uuid4 + +import httpx + +from configs import dify_config +from core.helper import ssrf_proxy +from extensions.ext_database import db +from extensions.ext_storage import storage +from models.model import MessageFile +from models.tools import ToolFile + +logger = logging.getLogger(__name__) + + +class ToolFileManager: + @staticmethod + def sign_file(tool_file_id: str, extension: str) -> str: + """ + sign file to get a temporary url + """ + base_url = dify_config.FILES_URL + file_preview_url = f"{base_url}/files/tools/{tool_file_id}{extension}" + + timestamp = str(int(time.time())) + nonce = os.urandom(16).hex() + data_to_sign = f"file-preview|{tool_file_id}|{timestamp}|{nonce}" + secret_key = dify_config.SECRET_KEY.encode() if dify_config.SECRET_KEY else b"" + sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest() + encoded_sign = base64.urlsafe_b64encode(sign).decode() + + return f"{file_preview_url}?timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}" + + @staticmethod + def verify_file(file_id: str, timestamp: str, nonce: str, sign: str) -> bool: + """ + verify signature + """ + data_to_sign = f"file-preview|{file_id}|{timestamp}|{nonce}" + secret_key = dify_config.SECRET_KEY.encode() if dify_config.SECRET_KEY else b"" + recalculated_sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest() + recalculated_encoded_sign = base64.urlsafe_b64encode(recalculated_sign).decode() + + # verify signature + if sign != recalculated_encoded_sign: + return False + + current_time = int(time.time()) + return current_time - int(timestamp) <= dify_config.FILES_ACCESS_TIMEOUT + + @staticmethod + def create_file_by_raw( + *, + user_id: str, + tenant_id: str, + conversation_id: Optional[str], + file_binary: bytes, + mimetype: str, + ) -> ToolFile: + extension = guess_extension(mimetype) or ".bin" + unique_name = uuid4().hex + filename = f"{unique_name}{extension}" + filepath = f"tools/{tenant_id}/{filename}" + storage.save(filepath, file_binary) + + tool_file = ToolFile( + user_id=user_id, + tenant_id=tenant_id, + conversation_id=conversation_id, + file_key=filepath, + mimetype=mimetype, + name=filename, + size=len(file_binary), + ) + + db.session.add(tool_file) + db.session.commit() + db.session.refresh(tool_file) + + return tool_file + + @staticmethod + def create_file_by_url( + user_id: str, + tenant_id: str, + conversation_id: str | None, + file_url: str, + ) -> ToolFile: + # try to download image + try: + response = ssrf_proxy.get(file_url) + response.raise_for_status() + blob = response.content + except httpx.TimeoutException as e: + raise ValueError(f"timeout when downloading file from {file_url}") + + mimetype = guess_type(file_url)[0] or "octet/stream" + extension = guess_extension(mimetype) or ".bin" + unique_name = uuid4().hex + filename = f"{unique_name}{extension}" + filepath = f"tools/{tenant_id}/{filename}" + storage.save(filepath, blob) + + tool_file = ToolFile( + user_id=user_id, + tenant_id=tenant_id, + conversation_id=conversation_id, + file_key=filepath, + mimetype=mimetype, + original_url=file_url, + name=filename, + size=len(blob), + ) + + db.session.add(tool_file) + db.session.commit() + + return tool_file + + @staticmethod + def get_file_binary(id: str) -> Union[tuple[bytes, str], None]: + """ + get file binary + + :param id: the id of the file + + :return: the binary of the file, mime type + """ + tool_file = ( + db.session.query(ToolFile) + .filter( + ToolFile.id == id, + ) + .first() + ) + + if not tool_file: + return None + + blob = storage.load_once(tool_file.file_key) + + return blob, tool_file.mimetype + + @staticmethod + def get_file_binary_by_message_file_id(id: str) -> Union[tuple[bytes, str], None]: + """ + get file binary + + :param id: the id of the file + + :return: the binary of the file, mime type + """ + message_file = ( + db.session.query(MessageFile) + .filter( + MessageFile.id == id, + ) + .first() + ) + + # Check if message_file is not None + if message_file is not None: + # get tool file id + if message_file.url is not None: + tool_file_id = message_file.url.split("/")[-1] + # trim extension + tool_file_id = tool_file_id.split(".")[0] + else: + tool_file_id = None + else: + tool_file_id = None + + tool_file = ( + db.session.query(ToolFile) + .filter( + ToolFile.id == tool_file_id, + ) + .first() + ) + + if not tool_file: + return None + + blob = storage.load_once(tool_file.file_key) + + return blob, tool_file.mimetype + + @staticmethod + def get_file_generator_by_tool_file_id(tool_file_id: str): + """ + get file binary + + :param tool_file_id: the id of the tool file + + :return: the binary of the file, mime type + """ + tool_file = ( + db.session.query(ToolFile) + .filter( + ToolFile.id == tool_file_id, + ) + .first() + ) + + if not tool_file: + return None, None + + stream = storage.load_stream(tool_file.file_key) + + return stream, tool_file + + +# init tool_file_parser +from core.file.tool_file_parser import tool_file_manager + +tool_file_manager["manager"] = ToolFileManager diff --git a/api/core/tools/tool_label_manager.py b/api/core/tools/tool_label_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..e53985951b0627ade3c01227ad65ad4df2e00542 --- /dev/null +++ b/api/core/tools/tool_label_manager.py @@ -0,0 +1,102 @@ +from core.tools.entities.values import default_tool_label_name_list +from core.tools.provider.api_tool_provider import ApiToolProviderController +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController +from core.tools.provider.tool_provider import ToolProviderController +from core.tools.provider.workflow_tool_provider import WorkflowToolProviderController +from extensions.ext_database import db +from models.tools import ToolLabelBinding + + +class ToolLabelManager: + @classmethod + def filter_tool_labels(cls, tool_labels: list[str]) -> list[str]: + """ + Filter tool labels + """ + tool_labels = [label for label in tool_labels if label in default_tool_label_name_list] + return list(set(tool_labels)) + + @classmethod + def update_tool_labels(cls, controller: ToolProviderController, labels: list[str]): + """ + Update tool labels + """ + labels = cls.filter_tool_labels(labels) + + if isinstance(controller, ApiToolProviderController | WorkflowToolProviderController): + provider_id = controller.provider_id + else: + raise ValueError("Unsupported tool type") + + # delete old labels + db.session.query(ToolLabelBinding).filter(ToolLabelBinding.tool_id == provider_id).delete() + + # insert new labels + for label in labels: + db.session.add( + ToolLabelBinding( + tool_id=provider_id, + tool_type=controller.provider_type.value, + label_name=label, + ) + ) + + db.session.commit() + + @classmethod + def get_tool_labels(cls, controller: ToolProviderController) -> list[str]: + """ + Get tool labels + """ + if isinstance(controller, ApiToolProviderController | WorkflowToolProviderController): + provider_id = controller.provider_id + elif isinstance(controller, BuiltinToolProviderController): + return controller.tool_labels + else: + raise ValueError("Unsupported tool type") + + labels: list[ToolLabelBinding] = ( + db.session.query(ToolLabelBinding.label_name) + .filter( + ToolLabelBinding.tool_id == provider_id, + ToolLabelBinding.tool_type == controller.provider_type.value, + ) + .all() + ) + + return [label.label_name for label in labels] + + @classmethod + def get_tools_labels(cls, tool_providers: list[ToolProviderController]) -> dict[str, list[str]]: + """ + Get tools labels + + :param tool_providers: list of tool providers + + :return: dict of tool labels + :key: tool id + :value: list of tool labels + """ + if not tool_providers: + return {} + + for controller in tool_providers: + if not isinstance(controller, ApiToolProviderController | WorkflowToolProviderController): + raise ValueError("Unsupported tool type") + + provider_ids = [ + controller.provider_id + for controller in tool_providers + if isinstance(controller, (ApiToolProviderController, WorkflowToolProviderController)) + ] + + labels: list[ToolLabelBinding] = ( + db.session.query(ToolLabelBinding).filter(ToolLabelBinding.tool_id.in_(provider_ids)).all() + ) + + tool_labels: dict[str, list[str]] = {label.tool_id: [] for label in labels} + + for label in labels: + tool_labels[label.tool_id].append(label.label_name) + + return tool_labels diff --git a/api/core/tools/tool_manager.py b/api/core/tools/tool_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..5b2173a4d0ad69d4de140e58cf07a7e9a5f91bc2 --- /dev/null +++ b/api/core/tools/tool_manager.py @@ -0,0 +1,694 @@ +import json +import logging +import mimetypes +from collections.abc import Generator +from os import listdir, path +from threading import Lock, Thread +from typing import Any, Optional, Union, cast + +from configs import dify_config +from core.agent.entities import AgentToolEntity +from core.app.entities.app_invoke_entities import InvokeFrom +from core.helper.module_import_helper import load_single_subclass_from_source +from core.helper.position_helper import is_filtered +from core.model_runtime.utils.encoders import jsonable_encoder +from core.tools.entities.api_entities import UserToolProvider, UserToolProviderTypeLiteral +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_entities import ApiProviderAuthType, ToolInvokeFrom, ToolParameter +from core.tools.errors import ToolNotFoundError, ToolProviderNotFoundError +from core.tools.provider.api_tool_provider import ApiToolProviderController +from core.tools.provider.builtin._positions import BuiltinToolProviderSort +from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController +from core.tools.provider.tool_provider import ToolProviderController +from core.tools.provider.workflow_tool_provider import WorkflowToolProviderController +from core.tools.tool.api_tool import ApiTool +from core.tools.tool.builtin_tool import BuiltinTool +from core.tools.tool.tool import Tool +from core.tools.tool_label_manager import ToolLabelManager +from core.tools.utils.configuration import ToolConfigurationManager, ToolParameterConfigurationManager +from core.workflow.nodes.tool.entities import ToolEntity +from extensions.ext_database import db +from models.tools import ApiToolProvider, BuiltinToolProvider, WorkflowToolProvider +from services.tools.tools_transform_service import ToolTransformService + +logger = logging.getLogger(__name__) + + +class ToolManager: + _builtin_provider_lock = Lock() + _builtin_providers: dict[str, BuiltinToolProviderController] = {} + _builtin_providers_loaded = False + _builtin_tools_labels: dict[str, Union[I18nObject, None]] = {} + + @classmethod + def get_builtin_provider(cls, provider: str) -> BuiltinToolProviderController: + """ + get the builtin provider + + :param provider: the name of the provider + :return: the provider + """ + if len(cls._builtin_providers) == 0: + # init the builtin providers + cls.load_builtin_providers_cache() + + if provider not in cls._builtin_providers: + raise ToolProviderNotFoundError(f"builtin provider {provider} not found") + + return cls._builtin_providers[provider] + + @classmethod + def get_builtin_tool(cls, provider: str, tool_name: str) -> Union[BuiltinTool, Tool]: + """ + get the builtin tool + + :param provider: the name of the provider + :param tool_name: the name of the tool + + :return: the provider, the tool + """ + provider_controller = cls.get_builtin_provider(provider) + tool = provider_controller.get_tool(tool_name) + if tool is None: + raise ToolNotFoundError(f"tool {tool_name} not found") + + return tool + + @classmethod + def get_tool( + cls, provider_type: str, provider_id: str, tool_name: str, tenant_id: Optional[str] = None + ) -> Union[BuiltinTool, ApiTool, Tool]: + """ + get the tool + + :param provider_type: the type of the provider + :param provider_name: the name of the provider + :param tool_name: the name of the tool + + :return: the tool + """ + if provider_type == "builtin": + return cls.get_builtin_tool(provider_id, tool_name) + elif provider_type == "api": + if tenant_id is None: + raise ValueError("tenant id is required for api provider") + api_provider, _ = cls.get_api_provider_controller(tenant_id, provider_id) + return api_provider.get_tool(tool_name) + elif provider_type == "app": + raise NotImplementedError("app provider not implemented") + else: + raise ToolProviderNotFoundError(f"provider type {provider_type} not found") + + @classmethod + def get_tool_runtime( + cls, + provider_type: str, + provider_id: str, + tool_name: str, + tenant_id: str, + invoke_from: InvokeFrom = InvokeFrom.DEBUGGER, + tool_invoke_from: ToolInvokeFrom = ToolInvokeFrom.AGENT, + ) -> Union[BuiltinTool, ApiTool, Tool]: + """ + get the tool runtime + + :param provider_type: the type of the provider + :param provider_name: the name of the provider + :param tool_name: the name of the tool + + :return: the tool + """ + controller: Union[BuiltinToolProviderController, ApiToolProviderController, WorkflowToolProviderController] + if provider_type == "builtin": + builtin_tool = cls.get_builtin_tool(provider_id, tool_name) + + # check if the builtin tool need credentials + provider_controller = cls.get_builtin_provider(provider_id) + if not provider_controller.need_credentials: + return builtin_tool.fork_tool_runtime( + runtime={ + "tenant_id": tenant_id, + "credentials": {}, + "invoke_from": invoke_from, + "tool_invoke_from": tool_invoke_from, + } + ) + + # get credentials + builtin_provider: Optional[BuiltinToolProvider] = ( + db.session.query(BuiltinToolProvider) + .filter( + BuiltinToolProvider.tenant_id == tenant_id, + BuiltinToolProvider.provider == provider_id, + ) + .first() + ) + + if builtin_provider is None: + raise ToolProviderNotFoundError(f"builtin provider {provider_id} not found") + + # decrypt the credentials + credentials = builtin_provider.credentials + controller = cls.get_builtin_provider(provider_id) + tool_configuration = ToolConfigurationManager(tenant_id=tenant_id, provider_controller=controller) + + decrypted_credentials = tool_configuration.decrypt_tool_credentials(credentials) + + return builtin_tool.fork_tool_runtime( + runtime={ + "tenant_id": tenant_id, + "credentials": decrypted_credentials, + "runtime_parameters": {}, + "invoke_from": invoke_from, + "tool_invoke_from": tool_invoke_from, + } + ) + + elif provider_type == "api": + if tenant_id is None: + raise ValueError("tenant id is required for api provider") + + api_provider, credentials = cls.get_api_provider_controller(tenant_id, provider_id) + + # decrypt the credentials + tool_configuration = ToolConfigurationManager(tenant_id=tenant_id, provider_controller=api_provider) + decrypted_credentials = tool_configuration.decrypt_tool_credentials(credentials) + + return api_provider.get_tool(tool_name).fork_tool_runtime( + runtime={ + "tenant_id": tenant_id, + "credentials": decrypted_credentials, + "invoke_from": invoke_from, + "tool_invoke_from": tool_invoke_from, + } + ) + elif provider_type == "workflow": + workflow_provider: Optional[WorkflowToolProvider] = ( + db.session.query(WorkflowToolProvider) + .filter(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == provider_id) + .first() + ) + + if workflow_provider is None: + raise ToolProviderNotFoundError(f"workflow provider {provider_id} not found") + + controller = ToolTransformService.workflow_provider_to_controller(db_provider=workflow_provider) + controller_tools: Optional[list[Tool]] = controller.get_tools( + user_id="", tenant_id=workflow_provider.tenant_id + ) + if controller_tools is None or len(controller_tools) == 0: + raise ToolProviderNotFoundError(f"workflow provider {provider_id} not found") + + return controller_tools[0].fork_tool_runtime( + runtime={ + "tenant_id": tenant_id, + "credentials": {}, + "invoke_from": invoke_from, + "tool_invoke_from": tool_invoke_from, + } + ) + elif provider_type == "app": + raise NotImplementedError("app provider not implemented") + else: + raise ToolProviderNotFoundError(f"provider type {provider_type} not found") + + @classmethod + def _init_runtime_parameter(cls, parameter_rule: ToolParameter, parameters: dict): + """ + init runtime parameter + """ + parameter_value = parameters.get(parameter_rule.name) + if not parameter_value and parameter_value != 0: + # get default value + parameter_value = parameter_rule.default + if not parameter_value and parameter_rule.required: + raise ValueError(f"tool parameter {parameter_rule.name} not found in tool config") + + if parameter_rule.type == ToolParameter.ToolParameterType.SELECT: + # check if tool_parameter_config in options + options = [x.value for x in parameter_rule.options or []] + if parameter_value is not None and parameter_value not in options: + raise ValueError( + f"tool parameter {parameter_rule.name} value {parameter_value} not in options {options}" + ) + + return parameter_rule.type.cast_value(parameter_value) + + @classmethod + def get_agent_tool_runtime( + cls, tenant_id: str, app_id: str, agent_tool: AgentToolEntity, invoke_from: InvokeFrom = InvokeFrom.DEBUGGER + ) -> Tool: + """ + get the agent tool runtime + """ + tool_entity = cls.get_tool_runtime( + provider_type=agent_tool.provider_type, + provider_id=agent_tool.provider_id, + tool_name=agent_tool.tool_name, + tenant_id=tenant_id, + invoke_from=invoke_from, + tool_invoke_from=ToolInvokeFrom.AGENT, + ) + runtime_parameters = {} + parameters = tool_entity.get_all_runtime_parameters() + for parameter in parameters: + # check file types + if ( + parameter.type + in { + ToolParameter.ToolParameterType.SYSTEM_FILES, + ToolParameter.ToolParameterType.FILE, + ToolParameter.ToolParameterType.FILES, + } + and parameter.required + ): + raise ValueError(f"file type parameter {parameter.name} not supported in agent") + + if parameter.form == ToolParameter.ToolParameterForm.FORM: + # save tool parameter to tool entity memory + value = cls._init_runtime_parameter(parameter, agent_tool.tool_parameters) + runtime_parameters[parameter.name] = value + + # decrypt runtime parameters + encryption_manager = ToolParameterConfigurationManager( + tenant_id=tenant_id, + tool_runtime=tool_entity, + provider_name=agent_tool.provider_id, + provider_type=agent_tool.provider_type, + identity_id=f"AGENT.{app_id}", + ) + runtime_parameters = encryption_manager.decrypt_tool_parameters(runtime_parameters) + if tool_entity.runtime is None or tool_entity.runtime.runtime_parameters is None: + raise ValueError("runtime not found or runtime parameters not found") + + tool_entity.runtime.runtime_parameters.update(runtime_parameters) + return tool_entity + + @classmethod + def get_workflow_tool_runtime( + cls, + tenant_id: str, + app_id: str, + node_id: str, + workflow_tool: "ToolEntity", + invoke_from: InvokeFrom = InvokeFrom.DEBUGGER, + ) -> Tool: + """ + get the workflow tool runtime + """ + tool_entity = cls.get_tool_runtime( + provider_type=workflow_tool.provider_type, + provider_id=workflow_tool.provider_id, + tool_name=workflow_tool.tool_name, + tenant_id=tenant_id, + invoke_from=invoke_from, + tool_invoke_from=ToolInvokeFrom.WORKFLOW, + ) + runtime_parameters = {} + parameters = tool_entity.get_all_runtime_parameters() + + for parameter in parameters: + # save tool parameter to tool entity memory + if parameter.form == ToolParameter.ToolParameterForm.FORM: + value = cls._init_runtime_parameter(parameter, workflow_tool.tool_configurations) + runtime_parameters[parameter.name] = value + + # decrypt runtime parameters + encryption_manager = ToolParameterConfigurationManager( + tenant_id=tenant_id, + tool_runtime=tool_entity, + provider_name=workflow_tool.provider_id, + provider_type=workflow_tool.provider_type, + identity_id=f"WORKFLOW.{app_id}.{node_id}", + ) + + if runtime_parameters: + runtime_parameters = encryption_manager.decrypt_tool_parameters(runtime_parameters) + + if tool_entity.runtime is None or tool_entity.runtime.runtime_parameters is None: + raise ValueError("runtime not found or runtime parameters not found") + + tool_entity.runtime.runtime_parameters.update(runtime_parameters) + return tool_entity + + @classmethod + def get_builtin_provider_icon(cls, provider: str) -> tuple[str, str]: + """ + get the absolute path of the icon of the builtin provider + + :param provider: the name of the provider + + :return: the absolute path of the icon, the mime type of the icon + """ + # get provider + provider_controller = cls.get_builtin_provider(provider) + if provider_controller.identity is None: + raise ToolProviderNotFoundError(f"builtin provider {provider} not found") + + absolute_path = path.join( + path.dirname(path.realpath(__file__)), + "provider", + "builtin", + provider, + "_assets", + provider_controller.identity.icon, + ) + # check if the icon exists + if not path.exists(absolute_path): + raise ToolProviderNotFoundError(f"builtin provider {provider} icon not found") + + # get the mime type + mime_type, _ = mimetypes.guess_type(absolute_path) + mime_type = mime_type or "application/octet-stream" + + return absolute_path, mime_type + + @classmethod + def list_builtin_providers(cls) -> Generator[BuiltinToolProviderController, None, None]: + # use cache first + if cls._builtin_providers_loaded: + yield from list(cls._builtin_providers.values()) + return + + with cls._builtin_provider_lock: + if cls._builtin_providers_loaded: + yield from list(cls._builtin_providers.values()) + return + + yield from cls._list_builtin_providers() + + @classmethod + def _list_builtin_providers(cls) -> Generator[BuiltinToolProviderController, None, None]: + """ + list all the builtin providers + """ + for provider in listdir(path.join(path.dirname(path.realpath(__file__)), "provider", "builtin")): + if provider.startswith("__"): + continue + + if path.isdir(path.join(path.dirname(path.realpath(__file__)), "provider", "builtin", provider)): + if provider.startswith("__"): + continue + + # init provider + try: + provider_class = load_single_subclass_from_source( + module_name=f"core.tools.provider.builtin.{provider}.{provider}", + script_path=path.join( + path.dirname(path.realpath(__file__)), "provider", "builtin", provider, f"{provider}.py" + ), + parent_type=BuiltinToolProviderController, + ) + provider_controller: BuiltinToolProviderController = provider_class() + if provider_controller.identity is None: + continue + cls._builtin_providers[provider_controller.identity.name] = provider_controller + for tool in provider_controller.get_tools() or []: + if tool.identity is None: + continue + cls._builtin_tools_labels[tool.identity.name] = tool.identity.label + yield provider_controller + + except Exception as e: + logger.exception(f"load builtin provider {provider}") + continue + # set builtin providers loaded + cls._builtin_providers_loaded = True + + @classmethod + def load_builtin_providers_cache(cls): + for _ in cls.list_builtin_providers(): + pass + + @classmethod + def clear_builtin_providers_cache(cls): + cls._builtin_providers = {} + cls._builtin_providers_loaded = False + + @classmethod + def get_tool_label(cls, tool_name: str) -> Union[I18nObject, None]: + """ + get the tool label + + :param tool_name: the name of the tool + + :return: the label of the tool + """ + if len(cls._builtin_tools_labels) == 0: + # init the builtin providers + cls.load_builtin_providers_cache() + + if tool_name not in cls._builtin_tools_labels: + return None + + return cls._builtin_tools_labels[tool_name] + + @classmethod + def user_list_providers( + cls, user_id: str, tenant_id: str, typ: UserToolProviderTypeLiteral + ) -> list[UserToolProvider]: + result_providers: dict[str, UserToolProvider] = {} + + filters = [] + if not typ: + filters.extend(["builtin", "api", "workflow"]) + else: + filters.append(typ) + + if "builtin" in filters: + # get builtin providers + builtin_providers = cls.list_builtin_providers() + + # get db builtin providers + db_builtin_providers: list[BuiltinToolProvider] = ( + db.session.query(BuiltinToolProvider).filter(BuiltinToolProvider.tenant_id == tenant_id).all() + ) + + find_db_builtin_provider = lambda provider: next( + (x for x in db_builtin_providers if x.provider == provider), None + ) + + # append builtin providers + for provider in builtin_providers: + # handle include, exclude + if provider.identity is None: + continue + if is_filtered( + include_set=cast(set[str], dify_config.POSITION_TOOL_INCLUDES_SET), + exclude_set=cast(set[str], dify_config.POSITION_TOOL_EXCLUDES_SET), + data=provider, + name_func=lambda x: x.identity.name, + ): + continue + + user_provider = ToolTransformService.builtin_provider_to_user_provider( + provider_controller=provider, + db_provider=find_db_builtin_provider(provider.identity.name), + decrypt_credentials=False, + ) + + result_providers[provider.identity.name] = user_provider + + # get db api providers + + if "api" in filters: + db_api_providers: list[ApiToolProvider] = ( + db.session.query(ApiToolProvider).filter(ApiToolProvider.tenant_id == tenant_id).all() + ) + + api_provider_controllers: list[dict[str, Any]] = [ + {"provider": provider, "controller": ToolTransformService.api_provider_to_controller(provider)} + for provider in db_api_providers + ] + + # get labels + labels = ToolLabelManager.get_tools_labels([x["controller"] for x in api_provider_controllers]) + + for api_provider_controller in api_provider_controllers: + user_provider = ToolTransformService.api_provider_to_user_provider( + provider_controller=api_provider_controller["controller"], + db_provider=api_provider_controller["provider"], + decrypt_credentials=False, + labels=labels.get(api_provider_controller["controller"].provider_id, []), + ) + result_providers[f"api_provider.{user_provider.name}"] = user_provider + + if "workflow" in filters: + # get workflow providers + workflow_providers: list[WorkflowToolProvider] = ( + db.session.query(WorkflowToolProvider).filter(WorkflowToolProvider.tenant_id == tenant_id).all() + ) + + workflow_provider_controllers: list[WorkflowToolProviderController] = [] + for provider in workflow_providers: + try: + workflow_provider_controllers.append( + ToolTransformService.workflow_provider_to_controller(db_provider=provider) + ) + except Exception as e: + # app has been deleted + pass + + labels = ToolLabelManager.get_tools_labels( + [cast(ToolProviderController, controller) for controller in workflow_provider_controllers] + ) + + for provider_controller in workflow_provider_controllers: + user_provider = ToolTransformService.workflow_provider_to_user_provider( + provider_controller=provider_controller, + labels=labels.get(provider_controller.provider_id, []), + ) + result_providers[f"workflow_provider.{user_provider.name}"] = user_provider + + return BuiltinToolProviderSort.sort(list(result_providers.values())) + + @classmethod + def get_api_provider_controller( + cls, tenant_id: str, provider_id: str + ) -> tuple[ApiToolProviderController, dict[str, Any]]: + """ + get the api provider + + :param provider_name: the name of the provider + + :return: the provider controller, the credentials + """ + provider: Optional[ApiToolProvider] = ( + db.session.query(ApiToolProvider) + .filter( + ApiToolProvider.id == provider_id, + ApiToolProvider.tenant_id == tenant_id, + ) + .first() + ) + + if provider is None: + raise ToolProviderNotFoundError(f"api provider {provider_id} not found") + + controller = ApiToolProviderController.from_db( + provider, + ApiProviderAuthType.API_KEY if provider.credentials["auth_type"] == "api_key" else ApiProviderAuthType.NONE, + ) + controller.load_bundled_tools(provider.tools) + + return controller, provider.credentials + + @classmethod + def user_get_api_provider(cls, provider: str, tenant_id: str) -> dict: + """ + get api provider + """ + """ + get tool provider + """ + provider_name = provider + provider_tool: Optional[ApiToolProvider] = ( + db.session.query(ApiToolProvider) + .filter( + ApiToolProvider.tenant_id == tenant_id, + ApiToolProvider.name == provider, + ) + .first() + ) + + if provider_tool is None: + raise ValueError(f"you have not added provider {provider_name}") + + try: + credentials = json.loads(provider_tool.credentials_str) or {} + except: + credentials = {} + + # package tool provider controller + controller = ApiToolProviderController.from_db( + provider_tool, + ApiProviderAuthType.API_KEY if credentials["auth_type"] == "api_key" else ApiProviderAuthType.NONE, + ) + # init tool configuration + tool_configuration = ToolConfigurationManager(tenant_id=tenant_id, provider_controller=controller) + + decrypted_credentials = tool_configuration.decrypt_tool_credentials(credentials) + masked_credentials = tool_configuration.mask_tool_credentials(decrypted_credentials) + + try: + icon = json.loads(provider_tool.icon) + except: + icon = {"background": "#252525", "content": "\ud83d\ude01"} + + # add tool labels + labels = ToolLabelManager.get_tool_labels(controller) + + return cast( + dict, + jsonable_encoder( + { + "schema_type": provider_tool.schema_type, + "schema": provider_tool.schema, + "tools": provider_tool.tools, + "icon": icon, + "description": provider_tool.description, + "credentials": masked_credentials, + "privacy_policy": provider_tool.privacy_policy, + "custom_disclaimer": provider_tool.custom_disclaimer, + "labels": labels, + } + ), + ) + + @classmethod + def get_tool_icon(cls, tenant_id: str, provider_type: str, provider_id: str) -> Union[str, dict]: + """ + get the tool icon + + :param tenant_id: the id of the tenant + :param provider_type: the type of the provider + :param provider_id: the id of the provider + :return: + """ + provider_type = provider_type + provider_id = provider_id + provider: Optional[Union[BuiltinToolProvider, ApiToolProvider, WorkflowToolProvider]] = None + if provider_type == "builtin": + return ( + dify_config.CONSOLE_API_URL + + "/console/api/workspaces/current/tool-provider/builtin/" + + provider_id + + "/icon" + ) + elif provider_type == "api": + try: + provider = ( + db.session.query(ApiToolProvider) + .filter(ApiToolProvider.tenant_id == tenant_id, ApiToolProvider.id == provider_id) + .first() + ) + if provider is None: + raise ToolProviderNotFoundError(f"api provider {provider_id} not found") + icon = json.loads(provider.icon) + if isinstance(icon, (str, dict)): + return icon + return {"background": "#252525", "content": "\ud83d\ude01"} + except: + return {"background": "#252525", "content": "\ud83d\ude01"} + elif provider_type == "workflow": + provider = ( + db.session.query(WorkflowToolProvider) + .filter(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == provider_id) + .first() + ) + if provider is None: + raise ToolProviderNotFoundError(f"workflow provider {provider_id} not found") + + try: + icon = json.loads(provider.icon) + if isinstance(icon, (str, dict)): + return icon + return {"background": "#252525", "content": "\ud83d\ude01"} + except: + return {"background": "#252525", "content": "\ud83d\ude01"} + else: + raise ValueError(f"provider type {provider_type} not found") + + +# preload builtin tool providers +Thread(target=ToolManager.load_builtin_providers_cache, name="pre_load_builtin_providers_cache", daemon=True).start() diff --git a/api/core/tools/utils/__init__.py b/api/core/tools/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/tools/utils/configuration.py b/api/core/tools/utils/configuration.py new file mode 100644 index 0000000000000000000000000000000000000000..d7720928644701e8307fe248d0a5cd79cebf95b1 --- /dev/null +++ b/api/core/tools/utils/configuration.py @@ -0,0 +1,256 @@ +from copy import deepcopy +from typing import Any + +from pydantic import BaseModel + +from core.helper import encrypter +from core.helper.tool_parameter_cache import ToolParameterCache, ToolParameterCacheType +from core.helper.tool_provider_cache import ToolProviderCredentialsCache, ToolProviderCredentialsCacheType +from core.tools.entities.tool_entities import ( + ToolParameter, + ToolProviderCredentials, +) +from core.tools.provider.tool_provider import ToolProviderController +from core.tools.tool.tool import Tool + + +class ToolConfigurationManager(BaseModel): + tenant_id: str + provider_controller: ToolProviderController + + def _deep_copy(self, credentials: dict[str, str]) -> dict[str, str]: + """ + deep copy credentials + """ + return deepcopy(credentials) + + def encrypt_tool_credentials(self, credentials: dict[str, str]) -> dict[str, str]: + """ + encrypt tool credentials with tenant id + + return a deep copy of credentials with encrypted values + """ + credentials = self._deep_copy(credentials) + + # get fields need to be decrypted + fields = self.provider_controller.get_credentials_schema() + for field_name, field in fields.items(): + if field.type == ToolProviderCredentials.CredentialsType.SECRET_INPUT: + if field_name in credentials: + encrypted = encrypter.encrypt_token(self.tenant_id, credentials[field_name]) + credentials[field_name] = encrypted + + return credentials + + def mask_tool_credentials(self, credentials: dict[str, Any]) -> dict[str, Any]: + """ + mask tool credentials + + return a deep copy of credentials with masked values + """ + credentials = self._deep_copy(credentials) + + # get fields need to be decrypted + fields = self.provider_controller.get_credentials_schema() + for field_name, field in fields.items(): + if field.type == ToolProviderCredentials.CredentialsType.SECRET_INPUT: + if field_name in credentials: + if len(credentials[field_name]) > 6: + credentials[field_name] = ( + credentials[field_name][:2] + + "*" * (len(credentials[field_name]) - 4) + + credentials[field_name][-2:] + ) + else: + credentials[field_name] = "*" * len(credentials[field_name]) + + return credentials + + def decrypt_tool_credentials(self, credentials: dict[str, str]) -> dict[str, str]: + """ + decrypt tool credentials with tenant id + + return a deep copy of credentials with decrypted values + """ + identity_id = "" + if self.provider_controller.identity: + identity_id = f"{self.provider_controller.provider_type.value}.{self.provider_controller.identity.name}" + + cache = ToolProviderCredentialsCache( + tenant_id=self.tenant_id, + identity_id=identity_id, + cache_type=ToolProviderCredentialsCacheType.PROVIDER, + ) + cached_credentials = cache.get() + if cached_credentials: + return cached_credentials + credentials = self._deep_copy(credentials) + # get fields need to be decrypted + fields = self.provider_controller.get_credentials_schema() + for field_name, field in fields.items(): + if field.type == ToolProviderCredentials.CredentialsType.SECRET_INPUT: + if field_name in credentials: + try: + credentials[field_name] = encrypter.decrypt_token(self.tenant_id, credentials[field_name]) + except: + pass + + cache.set(credentials) + return credentials + + def delete_tool_credentials_cache(self): + identity_id = "" + if self.provider_controller.identity: + identity_id = f"{self.provider_controller.provider_type.value}.{self.provider_controller.identity.name}" + + cache = ToolProviderCredentialsCache( + tenant_id=self.tenant_id, + identity_id=identity_id, + cache_type=ToolProviderCredentialsCacheType.PROVIDER, + ) + cache.delete() + + +class ToolParameterConfigurationManager(BaseModel): + """ + Tool parameter configuration manager + """ + + tenant_id: str + tool_runtime: Tool + provider_name: str + provider_type: str + identity_id: str + + def _deep_copy(self, parameters: dict[str, Any]) -> dict[str, Any]: + """ + deep copy parameters + """ + return deepcopy(parameters) + + def _merge_parameters(self) -> list[ToolParameter]: + """ + merge parameters + """ + # get tool parameters + tool_parameters = self.tool_runtime.parameters or [] + # get tool runtime parameters + runtime_parameters = self.tool_runtime.get_runtime_parameters() + # override parameters + current_parameters = tool_parameters.copy() + for runtime_parameter in runtime_parameters: + found = False + for index, parameter in enumerate(current_parameters): + if parameter.name == runtime_parameter.name and parameter.form == runtime_parameter.form: + current_parameters[index] = runtime_parameter + found = True + break + + if not found and runtime_parameter.form == ToolParameter.ToolParameterForm.FORM: + current_parameters.append(runtime_parameter) + + return current_parameters + + def mask_tool_parameters(self, parameters: dict[str, Any]) -> dict[str, Any]: + """ + mask tool parameters + + return a deep copy of parameters with masked values + """ + parameters = self._deep_copy(parameters) + + # override parameters + current_parameters = self._merge_parameters() + + for parameter in current_parameters: + if ( + parameter.form == ToolParameter.ToolParameterForm.FORM + and parameter.type == ToolParameter.ToolParameterType.SECRET_INPUT + ): + if parameter.name in parameters: + if len(parameters[parameter.name]) > 6: + parameters[parameter.name] = ( + parameters[parameter.name][:2] + + "*" * (len(parameters[parameter.name]) - 4) + + parameters[parameter.name][-2:] + ) + else: + parameters[parameter.name] = "*" * len(parameters[parameter.name]) + + return parameters + + def encrypt_tool_parameters(self, parameters: dict[str, Any]) -> dict[str, Any]: + """ + encrypt tool parameters with tenant id + + return a deep copy of parameters with encrypted values + """ + # override parameters + current_parameters = self._merge_parameters() + + parameters = self._deep_copy(parameters) + + for parameter in current_parameters: + if ( + parameter.form == ToolParameter.ToolParameterForm.FORM + and parameter.type == ToolParameter.ToolParameterType.SECRET_INPUT + ): + if parameter.name in parameters: + encrypted = encrypter.encrypt_token(self.tenant_id, parameters[parameter.name]) + parameters[parameter.name] = encrypted + + return parameters + + def decrypt_tool_parameters(self, parameters: dict[str, Any]) -> dict[str, Any]: + """ + decrypt tool parameters with tenant id + + return a deep copy of parameters with decrypted values + """ + if self.tool_runtime is None or self.tool_runtime.identity is None: + raise ValueError("tool_runtime is required") + + cache = ToolParameterCache( + tenant_id=self.tenant_id, + provider=f"{self.provider_type}.{self.provider_name}", + tool_name=self.tool_runtime.identity.name, + cache_type=ToolParameterCacheType.PARAMETER, + identity_id=self.identity_id, + ) + cached_parameters = cache.get() + if cached_parameters: + return cached_parameters + + # override parameters + current_parameters = self._merge_parameters() + has_secret_input = False + + for parameter in current_parameters: + if ( + parameter.form == ToolParameter.ToolParameterForm.FORM + and parameter.type == ToolParameter.ToolParameterType.SECRET_INPUT + ): + if parameter.name in parameters: + try: + has_secret_input = True + parameters[parameter.name] = encrypter.decrypt_token(self.tenant_id, parameters[parameter.name]) + except: + pass + + if has_secret_input: + cache.set(parameters) + + return parameters + + def delete_tool_parameters_cache(self): + if self.tool_runtime is None or self.tool_runtime.identity is None: + raise ValueError("tool_runtime is required") + + cache = ToolParameterCache( + tenant_id=self.tenant_id, + provider=f"{self.provider_type}.{self.provider_name}", + tool_name=self.tool_runtime.identity.name, + cache_type=ToolParameterCacheType.PARAMETER, + identity_id=self.identity_id, + ) + cache.delete() diff --git a/api/core/tools/utils/feishu_api_utils.py b/api/core/tools/utils/feishu_api_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ecf60045aa8dc5050323b631c0ca8b08c23f5302 --- /dev/null +++ b/api/core/tools/utils/feishu_api_utils.py @@ -0,0 +1,919 @@ +import json +from typing import Any, Optional, cast + +import httpx + +from core.tools.errors import ToolProviderCredentialValidationError +from extensions.ext_redis import redis_client + + +def auth(credentials): + app_id = credentials.get("app_id") + app_secret = credentials.get("app_secret") + if not app_id or not app_secret: + raise ToolProviderCredentialValidationError("app_id and app_secret is required") + try: + assert FeishuRequest(app_id, app_secret).tenant_access_token is not None + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) + + +def convert_add_records(json_str): + try: + data = json.loads(json_str) + if not isinstance(data, list): + raise ValueError("Parsed data must be a list") + converted_data = [{"fields": json.dumps(item, ensure_ascii=False)} for item in data] + return converted_data + except json.JSONDecodeError: + raise ValueError("The input string is not valid JSON") + except Exception as e: + raise ValueError(f"An error occurred while processing the data: {e}") + + +def convert_update_records(json_str): + try: + data = json.loads(json_str) + if not isinstance(data, list): + raise ValueError("Parsed data must be a list") + + converted_data = [ + {"fields": json.dumps(record["fields"], ensure_ascii=False), "record_id": record["record_id"]} + for record in data + if "fields" in record and "record_id" in record + ] + + if len(converted_data) != len(data): + raise ValueError("Each record must contain 'fields' and 'record_id'") + + return converted_data + except json.JSONDecodeError: + raise ValueError("The input string is not valid JSON") + except Exception as e: + raise ValueError(f"An error occurred while processing the data: {e}") + + +class FeishuRequest: + API_BASE_URL = "https://lark-plugin-api.solutionsuite.cn/lark-plugin" + + def __init__(self, app_id: str, app_secret: str): + self.app_id = app_id + self.app_secret = app_secret + + @property + def tenant_access_token(self): + feishu_tenant_access_token = f"tools:{self.app_id}:feishu_tenant_access_token" + if redis_client.exists(feishu_tenant_access_token): + return redis_client.get(feishu_tenant_access_token).decode() + res = self.get_tenant_access_token(self.app_id, self.app_secret) + redis_client.setex(feishu_tenant_access_token, res.get("expire"), res.get("tenant_access_token")) + return res.get("tenant_access_token") + + def _send_request( + self, + url: str, + method: str = "post", + require_token: bool = True, + payload: Optional[dict] = None, + params: Optional[dict] = None, + ): + headers = { + "Content-Type": "application/json", + "user-agent": "Dify", + } + if require_token: + headers["tenant-access-token"] = f"{self.tenant_access_token}" + res = httpx.request(method=method, url=url, headers=headers, json=payload, params=params, timeout=30).json() + if res.get("code") != 0: + raise Exception(res) + return res + + def get_tenant_access_token(self, app_id: str, app_secret: str) -> dict: + """ + API url: https://open.feishu.cn/document/server-docs/authentication-management/access-token/tenant_access_token_internal + Example Response: + { + "code": 0, + "msg": "ok", + "tenant_access_token": "t-caecc734c2e3328a62489fe0648c4b98779515d3", + "expire": 7200 + } + """ + url = f"{self.API_BASE_URL}/access_token/get_tenant_access_token" + payload = {"app_id": app_id, "app_secret": app_secret} + res: dict = self._send_request(url, require_token=False, payload=payload) + return res + + def create_document(self, title: str, content: str, folder_token: str) -> dict: + """ + API url: https://open.larkoffice.com/document/server-docs/docs/docs/docx-v1/document/create + Example Response: + { + "data": { + "title": "title", + "url": "https://svi136aogf123.feishu.cn/docx/VWbvd4fEdoW0WSxaY1McQTz8n7d", + "type": "docx", + "token": "VWbvd4fEdoW0WSxaY1McQTz8n7d" + }, + "log_id": "021721281231575fdbddc0200ff00060a9258ec0000103df61b5d", + "code": 0, + "msg": "创建飞书文档成功,请查看" + } + """ + url = f"{self.API_BASE_URL}/document/create_document" + payload = { + "title": title, + "content": content, + "folder_token": folder_token, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def write_document(self, document_id: str, content: str, position: str = "end") -> dict: + url = f"{self.API_BASE_URL}/document/write_document" + payload = {"document_id": document_id, "content": content, "position": position} + res: dict = self._send_request(url, payload=payload) + return res + + def get_document_content(self, document_id: str, mode: str = "markdown", lang: str = "0") -> str: + """ + API url: https://open.larkoffice.com/document/server-docs/docs/docs/docx-v1/document/raw_content + Example Response: + { + "code": 0, + "msg": "success", + "data": { + "content": "云文档\n多人实时协同,插入一切元素。不仅是在线文档,更是强大的创作和互动工具\n云文档:专为协作而生\n" + } + } + """ # noqa: E501 + params = { + "document_id": document_id, + "mode": mode, + "lang": lang, + } + url = f"{self.API_BASE_URL}/document/get_document_content" + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + return cast(str, res.get("data", {}).get("content")) + return "" + + def list_document_blocks( + self, document_id: str, page_token: str, user_id_type: str = "open_id", page_size: int = 500 + ) -> dict: + """ + API url: https://open.larkoffice.com/document/server-docs/docs/docs/docx-v1/document/list + """ + params = { + "user_id_type": user_id_type, + "document_id": document_id, + "page_size": page_size, + "page_token": page_token, + } + url = f"{self.API_BASE_URL}/document/list_document_blocks" + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def send_bot_message(self, receive_id_type: str, receive_id: str, msg_type: str, content: str) -> dict: + """ + API url: https://open.larkoffice.com/document/server-docs/im-v1/message/create + """ + url = f"{self.API_BASE_URL}/message/send_bot_message" + params = { + "receive_id_type": receive_id_type, + } + payload = { + "receive_id": receive_id, + "msg_type": msg_type, + "content": content.strip('"').replace(r"\"", '"').replace(r"\\", "\\"), + } + res: dict = self._send_request(url, params=params, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def send_webhook_message(self, webhook: str, msg_type: str, content: str) -> dict: + url = f"{self.API_BASE_URL}/message/send_webhook_message" + payload = { + "webhook": webhook, + "msg_type": msg_type, + "content": content.strip('"').replace(r"\"", '"').replace(r"\\", "\\"), + } + res: dict = self._send_request(url, require_token=False, payload=payload) + return res + + def get_chat_messages( + self, + container_id: str, + start_time: str, + end_time: str, + page_token: str, + sort_type: str = "ByCreateTimeAsc", + page_size: int = 20, + ) -> dict: + """ + API url: https://open.larkoffice.com/document/server-docs/im-v1/message/list + """ + url = f"{self.API_BASE_URL}/message/get_chat_messages" + params = { + "container_id": container_id, + "start_time": start_time, + "end_time": end_time, + "sort_type": sort_type, + "page_token": page_token, + "page_size": page_size, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def get_thread_messages( + self, container_id: str, page_token: str, sort_type: str = "ByCreateTimeAsc", page_size: int = 20 + ) -> dict: + """ + API url: https://open.larkoffice.com/document/server-docs/im-v1/message/list + """ + url = f"{self.API_BASE_URL}/message/get_thread_messages" + params = { + "container_id": container_id, + "sort_type": sort_type, + "page_token": page_token, + "page_size": page_size, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def create_task(self, summary: str, start_time: str, end_time: str, completed_time: str, description: str) -> dict: + # 创建任务 + url = f"{self.API_BASE_URL}/task/create_task" + payload = { + "summary": summary, + "start_time": start_time, + "end_time": end_time, + "completed_at": completed_time, + "description": description, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def update_task( + self, task_guid: str, summary: str, start_time: str, end_time: str, completed_time: str, description: str + ) -> dict: + # 更新任务 + url = f"{self.API_BASE_URL}/task/update_task" + payload = { + "task_guid": task_guid, + "summary": summary, + "start_time": start_time, + "end_time": end_time, + "completed_time": completed_time, + "description": description, + } + res: dict = self._send_request(url, method="PATCH", payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def delete_task(self, task_guid: str) -> dict: + # 删除任务 + url = f"{self.API_BASE_URL}/task/delete_task" + payload = { + "task_guid": task_guid, + } + res: dict = self._send_request(url, method="DELETE", payload=payload) + return res + + def add_members(self, task_guid: str, member_phone_or_email: str, member_role: str) -> dict: + # 删除任务 + url = f"{self.API_BASE_URL}/task/add_members" + payload = { + "task_guid": task_guid, + "member_phone_or_email": member_phone_or_email, + "member_role": member_role, + } + res: dict = self._send_request(url, payload=payload) + return res + + def get_wiki_nodes(self, space_id: str, parent_node_token: str, page_token: str, page_size: int = 20) -> dict: + # 获取知识库全部子节点列表 + url = f"{self.API_BASE_URL}/wiki/get_wiki_nodes" + payload = { + "space_id": space_id, + "parent_node_token": parent_node_token, + "page_token": page_token, + "page_size": page_size, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def get_primary_calendar(self, user_id_type: str = "open_id") -> dict: + url = f"{self.API_BASE_URL}/calendar/get_primary_calendar" + params = { + "user_id_type": user_id_type, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def create_event( + self, + summary: str, + description: str, + start_time: str, + end_time: str, + attendee_ability: str, + need_notification: bool = True, + auto_record: bool = False, + ) -> dict: + url = f"{self.API_BASE_URL}/calendar/create_event" + payload = { + "summary": summary, + "description": description, + "need_notification": need_notification, + "start_time": start_time, + "end_time": end_time, + "auto_record": auto_record, + "attendee_ability": attendee_ability, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def update_event( + self, + event_id: str, + summary: str, + description: str, + need_notification: bool, + start_time: str, + end_time: str, + auto_record: bool, + ) -> dict: + url = f"{self.API_BASE_URL}/calendar/update_event/{event_id}" + payload: dict[str, Any] = {} + if summary: + payload["summary"] = summary + if description: + payload["description"] = description + if start_time: + payload["start_time"] = start_time + if end_time: + payload["end_time"] = end_time + if need_notification: + payload["need_notification"] = need_notification + if auto_record: + payload["auto_record"] = auto_record + res: dict = self._send_request(url, method="PATCH", payload=payload) + return res + + def delete_event(self, event_id: str, need_notification: bool = True) -> dict: + url = f"{self.API_BASE_URL}/calendar/delete_event/{event_id}" + params = { + "need_notification": need_notification, + } + res: dict = self._send_request(url, method="DELETE", params=params) + return res + + def list_events(self, start_time: str, end_time: str, page_token: str, page_size: int = 50) -> dict: + url = f"{self.API_BASE_URL}/calendar/list_events" + params = { + "start_time": start_time, + "end_time": end_time, + "page_token": page_token, + "page_size": page_size, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def search_events( + self, + query: str, + start_time: str, + end_time: str, + page_token: str, + user_id_type: str = "open_id", + page_size: int = 20, + ) -> dict: + url = f"{self.API_BASE_URL}/calendar/search_events" + payload = { + "query": query, + "start_time": start_time, + "end_time": end_time, + "page_token": page_token, + "user_id_type": user_id_type, + "page_size": page_size, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def add_event_attendees(self, event_id: str, attendee_phone_or_email: str, need_notification: bool = True) -> dict: + # 参加日程参会人 + url = f"{self.API_BASE_URL}/calendar/add_event_attendees" + payload = { + "event_id": event_id, + "attendee_phone_or_email": attendee_phone_or_email, + "need_notification": need_notification, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def create_spreadsheet( + self, + title: str, + folder_token: str, + ) -> dict: + # 创建电子表格 + url = f"{self.API_BASE_URL}/spreadsheet/create_spreadsheet" + payload = { + "title": title, + "folder_token": folder_token, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def get_spreadsheet( + self, + spreadsheet_token: str, + user_id_type: str = "open_id", + ) -> dict: + # 获取电子表格信息 + url = f"{self.API_BASE_URL}/spreadsheet/get_spreadsheet" + params = { + "spreadsheet_token": spreadsheet_token, + "user_id_type": user_id_type, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def list_spreadsheet_sheets( + self, + spreadsheet_token: str, + ) -> dict: + # 列出电子表格的所有工作表 + url = f"{self.API_BASE_URL}/spreadsheet/list_spreadsheet_sheets" + params = { + "spreadsheet_token": spreadsheet_token, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def add_rows( + self, + spreadsheet_token: str, + sheet_id: str, + sheet_name: str, + length: int, + values: str, + ) -> dict: + # 增加行,在工作表最后添加 + url = f"{self.API_BASE_URL}/spreadsheet/add_rows" + payload = { + "spreadsheet_token": spreadsheet_token, + "sheet_id": sheet_id, + "sheet_name": sheet_name, + "length": length, + "values": values, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def add_cols( + self, + spreadsheet_token: str, + sheet_id: str, + sheet_name: str, + length: int, + values: str, + ) -> dict: + # 增加列,在工作表最后添加 + url = f"{self.API_BASE_URL}/spreadsheet/add_cols" + payload = { + "spreadsheet_token": spreadsheet_token, + "sheet_id": sheet_id, + "sheet_name": sheet_name, + "length": length, + "values": values, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def read_rows( + self, + spreadsheet_token: str, + sheet_id: str, + sheet_name: str, + start_row: int, + num_rows: int, + user_id_type: str = "open_id", + ) -> dict: + # 读取工作表行数据 + url = f"{self.API_BASE_URL}/spreadsheet/read_rows" + params = { + "spreadsheet_token": spreadsheet_token, + "sheet_id": sheet_id, + "sheet_name": sheet_name, + "start_row": start_row, + "num_rows": num_rows, + "user_id_type": user_id_type, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def read_cols( + self, + spreadsheet_token: str, + sheet_id: str, + sheet_name: str, + start_col: int, + num_cols: int, + user_id_type: str = "open_id", + ) -> dict: + # 读取工作表列数据 + url = f"{self.API_BASE_URL}/spreadsheet/read_cols" + params = { + "spreadsheet_token": spreadsheet_token, + "sheet_id": sheet_id, + "sheet_name": sheet_name, + "start_col": start_col, + "num_cols": num_cols, + "user_id_type": user_id_type, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def read_table( + self, + spreadsheet_token: str, + sheet_id: str, + sheet_name: str, + num_range: str, + query: str, + user_id_type: str = "open_id", + ) -> dict: + # 自定义读取行列数据 + url = f"{self.API_BASE_URL}/spreadsheet/read_table" + params = { + "spreadsheet_token": spreadsheet_token, + "sheet_id": sheet_id, + "sheet_name": sheet_name, + "range": num_range, + "query": query, + "user_id_type": user_id_type, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def create_base( + self, + name: str, + folder_token: str, + ) -> dict: + # 创建多维表格 + url = f"{self.API_BASE_URL}/base/create_base" + payload = { + "name": name, + "folder_token": folder_token, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def add_records( + self, + app_token: str, + table_id: str, + table_name: str, + records: str, + user_id_type: str = "open_id", + ) -> dict: + # 新增多条记录 + url = f"{self.API_BASE_URL}/base/add_records" + params = { + "app_token": app_token, + "table_id": table_id, + "table_name": table_name, + "user_id_type": user_id_type, + } + payload = { + "records": convert_add_records(records), + } + res: dict = self._send_request(url, params=params, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def update_records( + self, + app_token: str, + table_id: str, + table_name: str, + records: str, + user_id_type: str, + ) -> dict: + # 更新多条记录 + url = f"{self.API_BASE_URL}/base/update_records" + params = { + "app_token": app_token, + "table_id": table_id, + "table_name": table_name, + "user_id_type": user_id_type, + } + payload = { + "records": convert_update_records(records), + } + res: dict = self._send_request(url, params=params, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def delete_records( + self, + app_token: str, + table_id: str, + table_name: str, + record_ids: str, + ) -> dict: + # 删除多条记录 + url = f"{self.API_BASE_URL}/base/delete_records" + params = { + "app_token": app_token, + "table_id": table_id, + "table_name": table_name, + } + if not record_ids: + record_id_list = [] + else: + try: + record_id_list = json.loads(record_ids) + except json.JSONDecodeError: + raise ValueError("The input string is not valid JSON") + payload = { + "records": record_id_list, + } + res: dict = self._send_request(url, params=params, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def search_record( + self, + app_token: str, + table_id: str, + table_name: str, + view_id: str, + field_names: str, + sort: str, + filters: str, + page_token: str, + automatic_fields: bool = False, + user_id_type: str = "open_id", + page_size: int = 20, + ) -> dict: + # 查询记录,单次最多查询 500 行记录。 + url = f"{self.API_BASE_URL}/base/search_record" + params = { + "app_token": app_token, + "table_id": table_id, + "table_name": table_name, + "user_id_type": user_id_type, + "page_token": page_token, + "page_size": page_size, + } + + if not field_names: + field_name_list = [] + else: + try: + field_name_list = json.loads(field_names) + except json.JSONDecodeError: + raise ValueError("The input string is not valid JSON") + + if not sort: + sort_list = [] + else: + try: + sort_list = json.loads(sort) + except json.JSONDecodeError: + raise ValueError("The input string is not valid JSON") + + if not filters: + filter_dict = {} + else: + try: + filter_dict = json.loads(filters) + except json.JSONDecodeError: + raise ValueError("The input string is not valid JSON") + + payload: dict[str, Any] = {} + + if view_id: + payload["view_id"] = view_id + if field_names: + payload["field_names"] = field_name_list + if sort: + payload["sort"] = sort_list + if filters: + payload["filter"] = filter_dict + if automatic_fields: + payload["automatic_fields"] = automatic_fields + res: dict = self._send_request(url, params=params, payload=payload) + + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def get_base_info( + self, + app_token: str, + ) -> dict: + # 获取多维表格元数据 + url = f"{self.API_BASE_URL}/base/get_base_info" + params = { + "app_token": app_token, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def create_table( + self, + app_token: str, + table_name: str, + default_view_name: str, + fields: str, + ) -> dict: + # 新增一个数据表 + url = f"{self.API_BASE_URL}/base/create_table" + params = { + "app_token": app_token, + } + if not fields: + fields_list = [] + else: + try: + fields_list = json.loads(fields) + except json.JSONDecodeError: + raise ValueError("The input string is not valid JSON") + payload = { + "name": table_name, + "fields": fields_list, + } + if default_view_name: + payload["default_view_name"] = default_view_name + res: dict = self._send_request(url, params=params, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def delete_tables( + self, + app_token: str, + table_ids: str, + table_names: str, + ) -> dict: + # 删除多个数据表 + url = f"{self.API_BASE_URL}/base/delete_tables" + params = { + "app_token": app_token, + } + if not table_ids: + table_id_list = [] + else: + try: + table_id_list = json.loads(table_ids) + except json.JSONDecodeError: + raise ValueError("The input string is not valid JSON") + + if not table_names: + table_name_list = [] + else: + try: + table_name_list = json.loads(table_names) + except json.JSONDecodeError: + raise ValueError("The input string is not valid JSON") + + payload = { + "table_ids": table_id_list, + "table_names": table_name_list, + } + + res: dict = self._send_request(url, params=params, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def list_tables( + self, + app_token: str, + page_token: str, + page_size: int = 20, + ) -> dict: + # 列出多维表格下的全部数据表 + url = f"{self.API_BASE_URL}/base/list_tables" + params = { + "app_token": app_token, + "page_token": page_token, + "page_size": page_size, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def read_records( + self, + app_token: str, + table_id: str, + table_name: str, + record_ids: str, + user_id_type: str = "open_id", + ) -> dict: + url = f"{self.API_BASE_URL}/base/read_records" + params = { + "app_token": app_token, + "table_id": table_id, + "table_name": table_name, + } + if not record_ids: + record_id_list = [] + else: + try: + record_id_list = json.loads(record_ids) + except json.JSONDecodeError: + raise ValueError("The input string is not valid JSON") + payload = { + "record_ids": record_id_list, + "user_id_type": user_id_type, + } + res: dict = self._send_request(url, method="GET", params=params, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res diff --git a/api/core/tools/utils/lark_api_utils.py b/api/core/tools/utils/lark_api_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..de394a39bf5a008b3bd2be3e92c9e03cf527b403 --- /dev/null +++ b/api/core/tools/utils/lark_api_utils.py @@ -0,0 +1,851 @@ +import json +from typing import Any, Optional, cast + +import httpx + +from core.tools.errors import ToolProviderCredentialValidationError +from extensions.ext_redis import redis_client + + +def lark_auth(credentials): + app_id = credentials.get("app_id") + app_secret = credentials.get("app_secret") + if not app_id or not app_secret: + raise ToolProviderCredentialValidationError("app_id and app_secret is required") + try: + assert LarkRequest(app_id, app_secret).tenant_access_token is not None + except Exception as e: + raise ToolProviderCredentialValidationError(str(e)) + + +class LarkRequest: + API_BASE_URL = "https://lark-plugin-api.solutionsuite.ai/lark-plugin" + + def __init__(self, app_id: str, app_secret: str): + self.app_id = app_id + self.app_secret = app_secret + + def convert_add_records(self, json_str): + try: + data = json.loads(json_str) + if not isinstance(data, list): + raise ValueError("Parsed data must be a list") + converted_data = [{"fields": json.dumps(item, ensure_ascii=False)} for item in data] + return converted_data + except json.JSONDecodeError: + raise ValueError("The input string is not valid JSON") + except Exception as e: + raise ValueError(f"An error occurred while processing the data: {e}") + + def convert_update_records(self, json_str): + try: + data = json.loads(json_str) + if not isinstance(data, list): + raise ValueError("Parsed data must be a list") + + converted_data = [ + {"fields": json.dumps(record["fields"], ensure_ascii=False), "record_id": record["record_id"]} + for record in data + if "fields" in record and "record_id" in record + ] + + if len(converted_data) != len(data): + raise ValueError("Each record must contain 'fields' and 'record_id'") + + return converted_data + except json.JSONDecodeError: + raise ValueError("The input string is not valid JSON") + except Exception as e: + raise ValueError(f"An error occurred while processing the data: {e}") + + @property + def tenant_access_token(self) -> str: + feishu_tenant_access_token = f"tools:{self.app_id}:feishu_tenant_access_token" + if redis_client.exists(feishu_tenant_access_token): + return str(redis_client.get(feishu_tenant_access_token).decode()) + res: dict[str, str] = self.get_tenant_access_token(self.app_id, self.app_secret) + redis_client.setex(feishu_tenant_access_token, res.get("expire"), res.get("tenant_access_token")) + return res.get("tenant_access_token", "") + + def _send_request( + self, + url: str, + method: str = "post", + require_token: bool = True, + payload: Optional[dict] = None, + params: Optional[dict] = None, + ): + headers = { + "Content-Type": "application/json", + "user-agent": "Dify", + } + if require_token: + headers["tenant-access-token"] = f"{self.tenant_access_token}" + res = httpx.request(method=method, url=url, headers=headers, json=payload, params=params, timeout=30).json() + if res.get("code") != 0: + raise Exception(res) + return res + + def get_tenant_access_token(self, app_id: str, app_secret: str) -> dict: + url = f"{self.API_BASE_URL}/access_token/get_tenant_access_token" + payload = {"app_id": app_id, "app_secret": app_secret} + res: dict = self._send_request(url, require_token=False, payload=payload) + return res + + def create_document(self, title: str, content: str, folder_token: str) -> dict: + url = f"{self.API_BASE_URL}/document/create_document" + payload = { + "title": title, + "content": content, + "folder_token": folder_token, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def write_document(self, document_id: str, content: str, position: str = "end") -> dict: + url = f"{self.API_BASE_URL}/document/write_document" + payload = {"document_id": document_id, "content": content, "position": position} + res: dict = self._send_request(url, payload=payload) + return res + + def get_document_content(self, document_id: str, mode: str = "markdown", lang: str = "0") -> str | dict: + params = { + "document_id": document_id, + "mode": mode, + "lang": lang, + } + url = f"{self.API_BASE_URL}/document/get_document_content" + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + return cast(dict, res.get("data", {}).get("content")) + return "" + + def list_document_blocks( + self, document_id: str, page_token: str, user_id_type: str = "open_id", page_size: int = 500 + ) -> dict: + params = { + "user_id_type": user_id_type, + "document_id": document_id, + "page_size": page_size, + "page_token": page_token, + } + url = f"{self.API_BASE_URL}/document/list_document_blocks" + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def send_bot_message(self, receive_id_type: str, receive_id: str, msg_type: str, content: str) -> dict: + url = f"{self.API_BASE_URL}/message/send_bot_message" + params = { + "receive_id_type": receive_id_type, + } + payload = { + "receive_id": receive_id, + "msg_type": msg_type, + "content": content.strip('"').replace(r"\"", '"').replace(r"\\", "\\"), + } + res: dict = self._send_request(url, params=params, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def send_webhook_message(self, webhook: str, msg_type: str, content: str) -> dict: + url = f"{self.API_BASE_URL}/message/send_webhook_message" + payload = { + "webhook": webhook, + "msg_type": msg_type, + "content": content.strip('"').replace(r"\"", '"').replace(r"\\", "\\"), + } + res: dict = self._send_request(url, require_token=False, payload=payload) + return res + + def get_chat_messages( + self, + container_id: str, + start_time: str, + end_time: str, + page_token: str, + sort_type: str = "ByCreateTimeAsc", + page_size: int = 20, + ) -> dict: + url = f"{self.API_BASE_URL}/message/get_chat_messages" + params = { + "container_id": container_id, + "start_time": start_time, + "end_time": end_time, + "sort_type": sort_type, + "page_token": page_token, + "page_size": page_size, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def get_thread_messages( + self, container_id: str, page_token: str, sort_type: str = "ByCreateTimeAsc", page_size: int = 20 + ) -> dict: + url = f"{self.API_BASE_URL}/message/get_thread_messages" + params = { + "container_id": container_id, + "sort_type": sort_type, + "page_token": page_token, + "page_size": page_size, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def create_task(self, summary: str, start_time: str, end_time: str, completed_time: str, description: str) -> dict: + url = f"{self.API_BASE_URL}/task/create_task" + payload = { + "summary": summary, + "start_time": start_time, + "end_time": end_time, + "completed_at": completed_time, + "description": description, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def update_task( + self, task_guid: str, summary: str, start_time: str, end_time: str, completed_time: str, description: str + ) -> dict: + url = f"{self.API_BASE_URL}/task/update_task" + payload = { + "task_guid": task_guid, + "summary": summary, + "start_time": start_time, + "end_time": end_time, + "completed_time": completed_time, + "description": description, + } + res: dict = self._send_request(url, method="PATCH", payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def delete_task(self, task_guid: str) -> dict: + url = f"{self.API_BASE_URL}/task/delete_task" + payload = { + "task_guid": task_guid, + } + res: dict = self._send_request(url, method="DELETE", payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def add_members(self, task_guid: str, member_phone_or_email: str, member_role: str) -> dict: + url = f"{self.API_BASE_URL}/task/add_members" + payload = { + "task_guid": task_guid, + "member_phone_or_email": member_phone_or_email, + "member_role": member_role, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def get_wiki_nodes(self, space_id: str, parent_node_token: str, page_token: str, page_size: int = 20) -> dict: + url = f"{self.API_BASE_URL}/wiki/get_wiki_nodes" + payload = { + "space_id": space_id, + "parent_node_token": parent_node_token, + "page_token": page_token, + "page_size": page_size, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def get_primary_calendar(self, user_id_type: str = "open_id") -> dict: + url = f"{self.API_BASE_URL}/calendar/get_primary_calendar" + params = { + "user_id_type": user_id_type, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def create_event( + self, + summary: str, + description: str, + start_time: str, + end_time: str, + attendee_ability: str, + need_notification: bool = True, + auto_record: bool = False, + ) -> dict: + url = f"{self.API_BASE_URL}/calendar/create_event" + payload = { + "summary": summary, + "description": description, + "need_notification": need_notification, + "start_time": start_time, + "end_time": end_time, + "auto_record": auto_record, + "attendee_ability": attendee_ability, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def update_event( + self, + event_id: str, + summary: str, + description: str, + need_notification: bool, + start_time: str, + end_time: str, + auto_record: bool, + ) -> dict: + url = f"{self.API_BASE_URL}/calendar/update_event/{event_id}" + payload: dict[str, Any] = {} + if summary: + payload["summary"] = summary + if description: + payload["description"] = description + if start_time: + payload["start_time"] = start_time + if end_time: + payload["end_time"] = end_time + if need_notification: + payload["need_notification"] = need_notification + if auto_record: + payload["auto_record"] = auto_record + res: dict = self._send_request(url, method="PATCH", payload=payload) + return res + + def delete_event(self, event_id: str, need_notification: bool = True) -> dict: + url = f"{self.API_BASE_URL}/calendar/delete_event/{event_id}" + params = { + "need_notification": need_notification, + } + res: dict = self._send_request(url, method="DELETE", params=params) + return res + + def list_events(self, start_time: str, end_time: str, page_token: str, page_size: int = 50) -> dict: + url = f"{self.API_BASE_URL}/calendar/list_events" + params = { + "start_time": start_time, + "end_time": end_time, + "page_token": page_token, + "page_size": page_size, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def search_events( + self, + query: str, + start_time: str, + end_time: str, + page_token: str, + user_id_type: str = "open_id", + page_size: int = 20, + ) -> dict: + url = f"{self.API_BASE_URL}/calendar/search_events" + payload = { + "query": query, + "start_time": start_time, + "end_time": end_time, + "page_token": page_token, + "user_id_type": user_id_type, + "page_size": page_size, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def add_event_attendees(self, event_id: str, attendee_phone_or_email: str, need_notification: bool = True) -> dict: + url = f"{self.API_BASE_URL}/calendar/add_event_attendees" + payload = { + "event_id": event_id, + "attendee_phone_or_email": attendee_phone_or_email, + "need_notification": need_notification, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def create_spreadsheet( + self, + title: str, + folder_token: str, + ) -> dict: + url = f"{self.API_BASE_URL}/spreadsheet/create_spreadsheet" + payload = { + "title": title, + "folder_token": folder_token, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def get_spreadsheet( + self, + spreadsheet_token: str, + user_id_type: str = "open_id", + ) -> dict: + url = f"{self.API_BASE_URL}/spreadsheet/get_spreadsheet" + params = { + "spreadsheet_token": spreadsheet_token, + "user_id_type": user_id_type, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def list_spreadsheet_sheets( + self, + spreadsheet_token: str, + ) -> dict: + url = f"{self.API_BASE_URL}/spreadsheet/list_spreadsheet_sheets" + params = { + "spreadsheet_token": spreadsheet_token, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def add_rows( + self, + spreadsheet_token: str, + sheet_id: str, + sheet_name: str, + length: int, + values: str, + ) -> dict: + url = f"{self.API_BASE_URL}/spreadsheet/add_rows" + payload = { + "spreadsheet_token": spreadsheet_token, + "sheet_id": sheet_id, + "sheet_name": sheet_name, + "length": length, + "values": values, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def add_cols( + self, + spreadsheet_token: str, + sheet_id: str, + sheet_name: str, + length: int, + values: str, + ) -> dict: + url = f"{self.API_BASE_URL}/spreadsheet/add_cols" + payload = { + "spreadsheet_token": spreadsheet_token, + "sheet_id": sheet_id, + "sheet_name": sheet_name, + "length": length, + "values": values, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def read_rows( + self, + spreadsheet_token: str, + sheet_id: str, + sheet_name: str, + start_row: int, + num_rows: int, + user_id_type: str = "open_id", + ) -> dict: + url = f"{self.API_BASE_URL}/spreadsheet/read_rows" + params = { + "spreadsheet_token": spreadsheet_token, + "sheet_id": sheet_id, + "sheet_name": sheet_name, + "start_row": start_row, + "num_rows": num_rows, + "user_id_type": user_id_type, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def read_cols( + self, + spreadsheet_token: str, + sheet_id: str, + sheet_name: str, + start_col: int, + num_cols: int, + user_id_type: str = "open_id", + ) -> dict: + url = f"{self.API_BASE_URL}/spreadsheet/read_cols" + params = { + "spreadsheet_token": spreadsheet_token, + "sheet_id": sheet_id, + "sheet_name": sheet_name, + "start_col": start_col, + "num_cols": num_cols, + "user_id_type": user_id_type, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def read_table( + self, + spreadsheet_token: str, + sheet_id: str, + sheet_name: str, + num_range: str, + query: str, + user_id_type: str = "open_id", + ) -> dict: + url = f"{self.API_BASE_URL}/spreadsheet/read_table" + params = { + "spreadsheet_token": spreadsheet_token, + "sheet_id": sheet_id, + "sheet_name": sheet_name, + "range": num_range, + "query": query, + "user_id_type": user_id_type, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def create_base( + self, + name: str, + folder_token: str, + ) -> dict: + url = f"{self.API_BASE_URL}/base/create_base" + payload = { + "name": name, + "folder_token": folder_token, + } + res: dict = self._send_request(url, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def add_records( + self, + app_token: str, + table_id: str, + table_name: str, + records: str, + user_id_type: str = "open_id", + ) -> dict: + url = f"{self.API_BASE_URL}/base/add_records" + params = { + "app_token": app_token, + "table_id": table_id, + "table_name": table_name, + "user_id_type": user_id_type, + } + payload = { + "records": self.convert_add_records(records), + } + res: dict = self._send_request(url, params=params, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def update_records( + self, + app_token: str, + table_id: str, + table_name: str, + records: str, + user_id_type: str, + ) -> dict: + url = f"{self.API_BASE_URL}/base/update_records" + params = { + "app_token": app_token, + "table_id": table_id, + "table_name": table_name, + "user_id_type": user_id_type, + } + payload = { + "records": self.convert_update_records(records), + } + res: dict = self._send_request(url, params=params, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def delete_records( + self, + app_token: str, + table_id: str, + table_name: str, + record_ids: str, + ) -> dict: + url = f"{self.API_BASE_URL}/base/delete_records" + params = { + "app_token": app_token, + "table_id": table_id, + "table_name": table_name, + } + if not record_ids: + record_id_list = [] + else: + try: + record_id_list = json.loads(record_ids) + except json.JSONDecodeError: + raise ValueError("The input string is not valid JSON") + payload = { + "records": record_id_list, + } + res: dict = self._send_request(url, params=params, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def search_record( + self, + app_token: str, + table_id: str, + table_name: str, + view_id: str, + field_names: str, + sort: str, + filters: str, + page_token: str, + automatic_fields: bool = False, + user_id_type: str = "open_id", + page_size: int = 20, + ) -> dict: + url = f"{self.API_BASE_URL}/base/search_record" + + params = { + "app_token": app_token, + "table_id": table_id, + "table_name": table_name, + "user_id_type": user_id_type, + "page_token": page_token, + "page_size": page_size, + } + + if not field_names: + field_name_list = [] + else: + try: + field_name_list = json.loads(field_names) + except json.JSONDecodeError: + raise ValueError("The input string is not valid JSON") + + if not sort: + sort_list = [] + else: + try: + sort_list = json.loads(sort) + except json.JSONDecodeError: + raise ValueError("The input string is not valid JSON") + + if not filters: + filter_dict = {} + else: + try: + filter_dict = json.loads(filters) + except json.JSONDecodeError: + raise ValueError("The input string is not valid JSON") + + payload: dict[str, Any] = {} + + if view_id: + payload["view_id"] = view_id + if field_names: + payload["field_names"] = field_name_list + if sort: + payload["sort"] = sort_list + if filters: + payload["filter"] = filter_dict + if automatic_fields: + payload["automatic_fields"] = automatic_fields + res: dict = self._send_request(url, params=params, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def get_base_info( + self, + app_token: str, + ) -> dict: + url = f"{self.API_BASE_URL}/base/get_base_info" + params = { + "app_token": app_token, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def create_table( + self, + app_token: str, + table_name: str, + default_view_name: str, + fields: str, + ) -> dict: + url = f"{self.API_BASE_URL}/base/create_table" + params = { + "app_token": app_token, + } + if not fields: + fields_list = [] + else: + try: + fields_list = json.loads(fields) + except json.JSONDecodeError: + raise ValueError("The input string is not valid JSON") + payload = { + "name": table_name, + "fields": fields_list, + } + if default_view_name: + payload["default_view_name"] = default_view_name + res: dict = self._send_request(url, params=params, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def delete_tables( + self, + app_token: str, + table_ids: str, + table_names: str, + ) -> dict: + url = f"{self.API_BASE_URL}/base/delete_tables" + params = { + "app_token": app_token, + } + if not table_ids: + table_id_list = [] + else: + try: + table_id_list = json.loads(table_ids) + except json.JSONDecodeError: + raise ValueError("The input string is not valid JSON") + + if not table_names: + table_name_list = [] + else: + try: + table_name_list = json.loads(table_names) + except json.JSONDecodeError: + raise ValueError("The input string is not valid JSON") + + payload = { + "table_ids": table_id_list, + "table_names": table_name_list, + } + res: dict = self._send_request(url, params=params, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def list_tables( + self, + app_token: str, + page_token: str, + page_size: int = 20, + ) -> dict: + url = f"{self.API_BASE_URL}/base/list_tables" + params = { + "app_token": app_token, + "page_token": page_token, + "page_size": page_size, + } + res: dict = self._send_request(url, method="GET", params=params) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res + + def read_records( + self, + app_token: str, + table_id: str, + table_name: str, + record_ids: str, + user_id_type: str = "open_id", + ) -> dict: + url = f"{self.API_BASE_URL}/base/read_records" + params = { + "app_token": app_token, + "table_id": table_id, + "table_name": table_name, + } + if not record_ids: + record_id_list = [] + else: + try: + record_id_list = json.loads(record_ids) + except json.JSONDecodeError: + raise ValueError("The input string is not valid JSON") + payload = { + "record_ids": record_id_list, + "user_id_type": user_id_type, + } + res: dict = self._send_request(url, method="POST", params=params, payload=payload) + if "data" in res: + data: dict = res.get("data", {}) + return data + return res diff --git a/api/core/tools/utils/message_transformer.py b/api/core/tools/utils/message_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..b28953264cfc7996157f15bb076698f25e3d2bc3 --- /dev/null +++ b/api/core/tools/utils/message_transformer.py @@ -0,0 +1,125 @@ +import logging +from mimetypes import guess_extension +from typing import Optional + +from core.file import File, FileTransferMethod, FileType +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.tools.tool_file_manager import ToolFileManager + +logger = logging.getLogger(__name__) + + +class ToolFileMessageTransformer: + @classmethod + def transform_tool_invoke_messages( + cls, messages: list[ToolInvokeMessage], user_id: str, tenant_id: str, conversation_id: str | None + ) -> list[ToolInvokeMessage]: + """ + Transform tool message and handle file download + """ + result = [] + + for message in messages: + if message.type in {ToolInvokeMessage.MessageType.TEXT, ToolInvokeMessage.MessageType.LINK}: + result.append(message) + elif message.type == ToolInvokeMessage.MessageType.IMAGE and isinstance(message.message, str): + # try to download image + try: + file = ToolFileManager.create_file_by_url( + user_id=user_id, tenant_id=tenant_id, conversation_id=conversation_id, file_url=message.message + ) + + url = f"/files/tools/{file.id}{guess_extension(file.mimetype) or '.png'}" + + result.append( + ToolInvokeMessage( + type=ToolInvokeMessage.MessageType.IMAGE_LINK, + message=url, + save_as=message.save_as, + meta=message.meta.copy() if message.meta is not None else {}, + ) + ) + except Exception as e: + logger.exception(f"Failed to download image from {url}") + result.append( + ToolInvokeMessage( + type=ToolInvokeMessage.MessageType.TEXT, + message=f"Failed to download image: {message.message}, please try to download it manually.", + meta=message.meta.copy() if message.meta is not None else {}, + save_as=message.save_as, + ) + ) + elif message.type == ToolInvokeMessage.MessageType.BLOB: + # get mime type and save blob to storage + assert message.meta is not None + mimetype = message.meta.get("mime_type", "octet/stream") + # if message is str, encode it to bytes + if isinstance(message.message, str): + message.message = message.message.encode("utf-8") + + # FIXME: should do a type check here. + assert isinstance(message.message, bytes) + file = ToolFileManager.create_file_by_raw( + user_id=user_id, + tenant_id=tenant_id, + conversation_id=conversation_id, + file_binary=message.message, + mimetype=mimetype, + ) + + url = cls.get_tool_file_url(tool_file_id=file.id, extension=guess_extension(file.mimetype)) + + # check if file is image + if "image" in mimetype: + result.append( + ToolInvokeMessage( + type=ToolInvokeMessage.MessageType.IMAGE_LINK, + message=url, + save_as=message.save_as, + meta=message.meta.copy() if message.meta is not None else {}, + ) + ) + else: + result.append( + ToolInvokeMessage( + type=ToolInvokeMessage.MessageType.LINK, + message=url, + save_as=message.save_as, + meta=message.meta.copy() if message.meta is not None else {}, + ) + ) + elif message.type == ToolInvokeMessage.MessageType.FILE: + assert message.meta is not None + file_mata = message.meta.get("file") + if isinstance(file_mata, File): + if file_mata.transfer_method == FileTransferMethod.TOOL_FILE: + assert file_mata.related_id is not None + url = cls.get_tool_file_url(tool_file_id=file_mata.related_id, extension=file_mata.extension) + if file_mata.type == FileType.IMAGE: + result.append( + ToolInvokeMessage( + type=ToolInvokeMessage.MessageType.IMAGE_LINK, + message=url, + save_as=message.save_as, + meta=message.meta.copy() if message.meta is not None else {}, + ) + ) + else: + result.append( + ToolInvokeMessage( + type=ToolInvokeMessage.MessageType.LINK, + message=url, + save_as=message.save_as, + meta=message.meta.copy() if message.meta is not None else {}, + ) + ) + else: + result.append(message) + else: + result.append(message) + + return result + + @classmethod + def get_tool_file_url(cls, tool_file_id: str, extension: Optional[str]) -> str: + return f"/files/tools/{tool_file_id}{extension or '.bin'}" diff --git a/api/core/tools/utils/model_invocation_utils.py b/api/core/tools/utils/model_invocation_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3689dcc9e5ebfdff244b34f96cbccbe951a646e7 --- /dev/null +++ b/api/core/tools/utils/model_invocation_utils.py @@ -0,0 +1,173 @@ +""" +For some reason, model will be used in tools like WebScraperTool, WikipediaSearchTool etc. + +Therefore, a model manager is needed to list/invoke/validate models. +""" + +import json +from typing import Optional, cast + +from core.model_manager import ModelManager +from core.model_runtime.entities.llm_entities import LLMResult +from core.model_runtime.entities.message_entities import PromptMessage +from core.model_runtime.entities.model_entities import ModelType +from core.model_runtime.errors.invoke import ( + InvokeAuthorizationError, + InvokeBadRequestError, + InvokeConnectionError, + InvokeRateLimitError, + InvokeServerUnavailableError, +) +from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel, ModelPropertyKey +from core.model_runtime.utils.encoders import jsonable_encoder +from extensions.ext_database import db +from models.tools import ToolModelInvoke + + +class InvokeModelError(Exception): + pass + + +class ModelInvocationUtils: + @staticmethod + def get_max_llm_context_tokens( + tenant_id: str, + ) -> int: + """ + get max llm context tokens of the model + """ + model_manager = ModelManager() + model_instance = model_manager.get_default_model_instance( + tenant_id=tenant_id, + model_type=ModelType.LLM, + ) + + if not model_instance: + raise InvokeModelError("Model not found") + + llm_model = cast(LargeLanguageModel, model_instance.model_type_instance) + schema = llm_model.get_model_schema(model_instance.model, model_instance.credentials) + + if not schema: + raise InvokeModelError("No model schema found") + + max_tokens: Optional[int] = schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE, None) + if max_tokens is None: + return 2048 + + return max_tokens + + @staticmethod + def calculate_tokens(tenant_id: str, prompt_messages: list[PromptMessage]) -> int: + """ + calculate tokens from prompt messages and model parameters + """ + + # get model instance + model_manager = ModelManager() + model_instance = model_manager.get_default_model_instance(tenant_id=tenant_id, model_type=ModelType.LLM) + + if not model_instance: + raise InvokeModelError("Model not found") + + # get tokens + tokens = model_instance.get_llm_num_tokens(prompt_messages) + + return tokens + + @staticmethod + def invoke( + user_id: str, tenant_id: str, tool_type: str, tool_name: str, prompt_messages: list[PromptMessage] + ) -> LLMResult: + """ + invoke model with parameters in user's own context + + :param user_id: user id + :param tenant_id: tenant id, the tenant id of the creator of the tool + :param tool_provider: tool provider + :param tool_id: tool id + :param tool_name: tool name + :param provider: model provider + :param model: model name + :param model_parameters: model parameters + :param prompt_messages: prompt messages + :return: AssistantPromptMessage + """ + + # get model manager + model_manager = ModelManager() + # get model instance + model_instance = model_manager.get_default_model_instance( + tenant_id=tenant_id, + model_type=ModelType.LLM, + ) + + # get prompt tokens + prompt_tokens = model_instance.get_llm_num_tokens(prompt_messages) + + model_parameters = { + "temperature": 0.8, + "top_p": 0.8, + } + + # create tool model invoke + tool_model_invoke = ToolModelInvoke( + user_id=user_id, + tenant_id=tenant_id, + provider=model_instance.provider, + tool_type=tool_type, + tool_name=tool_name, + model_parameters=json.dumps(model_parameters), + prompt_messages=json.dumps(jsonable_encoder(prompt_messages)), + model_response="", + prompt_tokens=prompt_tokens, + answer_tokens=0, + answer_unit_price=0, + answer_price_unit=0, + provider_response_latency=0, + total_price=0, + currency="USD", + ) + + db.session.add(tool_model_invoke) + db.session.commit() + + try: + response: LLMResult = cast( + LLMResult, + model_instance.invoke_llm( + prompt_messages=prompt_messages, + model_parameters=model_parameters, + tools=[], + stop=[], + stream=False, + user=user_id, + callbacks=[], + ), + ) + except InvokeRateLimitError as e: + raise InvokeModelError(f"Invoke rate limit error: {e}") + except InvokeBadRequestError as e: + raise InvokeModelError(f"Invoke bad request error: {e}") + except InvokeConnectionError as e: + raise InvokeModelError(f"Invoke connection error: {e}") + except InvokeAuthorizationError as e: + raise InvokeModelError("Invoke authorization error") + except InvokeServerUnavailableError as e: + raise InvokeModelError(f"Invoke server unavailable error: {e}") + except Exception as e: + raise InvokeModelError(f"Invoke error: {e}") + + # update tool model invoke + tool_model_invoke.model_response = response.message.content + if response.usage: + tool_model_invoke.answer_tokens = response.usage.completion_tokens + tool_model_invoke.answer_unit_price = response.usage.completion_unit_price + tool_model_invoke.answer_price_unit = response.usage.completion_price_unit + tool_model_invoke.provider_response_latency = response.usage.latency + tool_model_invoke.total_price = response.usage.total_price + tool_model_invoke.currency = response.usage.currency + + db.session.commit() + + return response diff --git a/api/core/tools/utils/parser.py b/api/core/tools/utils/parser.py new file mode 100644 index 0000000000000000000000000000000000000000..b15a86b5c0a5cfd5e6e7d8d64cb7119cd4c7d637 --- /dev/null +++ b/api/core/tools/utils/parser.py @@ -0,0 +1,379 @@ +import re +import uuid +from json import dumps as json_dumps +from json import loads as json_loads +from json.decoder import JSONDecodeError +from typing import Optional + +from flask import request +from requests import get +from yaml import YAMLError, safe_load # type: ignore + +from core.tools.entities.common_entities import I18nObject +from core.tools.entities.tool_bundle import ApiToolBundle +from core.tools.entities.tool_entities import ApiProviderSchemaType, ToolParameter +from core.tools.errors import ToolApiSchemaError, ToolNotSupportedError, ToolProviderNotFoundError + + +class ApiBasedToolSchemaParser: + @staticmethod + def parse_openapi_to_tool_bundle( + openapi: dict, extra_info: Optional[dict], warning: Optional[dict] + ) -> list[ApiToolBundle]: + warning = warning if warning is not None else {} + extra_info = extra_info if extra_info is not None else {} + + # set description to extra_info + extra_info["description"] = openapi["info"].get("description", "") + + if len(openapi["servers"]) == 0: + raise ToolProviderNotFoundError("No server found in the openapi yaml.") + + server_url = openapi["servers"][0]["url"] + request_env = request.headers.get("X-Request-Env") + if request_env: + matched_servers = [server["url"] for server in openapi["servers"] if server["env"] == request_env] + server_url = matched_servers[0] if matched_servers else server_url + + # list all interfaces + interfaces = [] + for path, path_item in openapi["paths"].items(): + methods = ["get", "post", "put", "delete", "patch", "head", "options", "trace"] + for method in methods: + if method in path_item: + interfaces.append( + { + "path": path, + "method": method, + "operation": path_item[method], + } + ) + + # get all parameters + bundles = [] + for interface in interfaces: + # convert parameters + parameters = [] + if "parameters" in interface["operation"]: + for parameter in interface["operation"]["parameters"]: + tool_parameter = ToolParameter( + name=parameter["name"], + label=I18nObject(en_US=parameter["name"], zh_Hans=parameter["name"]), + human_description=I18nObject( + en_US=parameter.get("description", ""), zh_Hans=parameter.get("description", "") + ), + type=ToolParameter.ToolParameterType.STRING, + required=parameter.get("required", False), + form=ToolParameter.ToolParameterForm.LLM, + llm_description=parameter.get("description"), + default=parameter["schema"]["default"] + if "schema" in parameter and "default" in parameter["schema"] + else None, + placeholder=I18nObject( + en_US=parameter.get("description", ""), zh_Hans=parameter.get("description", "") + ), + ) + + # check if there is a type + typ = ApiBasedToolSchemaParser._get_tool_parameter_type(parameter) + if typ: + tool_parameter.type = typ + + parameters.append(tool_parameter) + # create tool bundle + # check if there is a request body + if "requestBody" in interface["operation"]: + request_body = interface["operation"]["requestBody"] + if "content" in request_body: + for content_type, content in request_body["content"].items(): + # if there is a reference, get the reference and overwrite the content + if "schema" not in content: + continue + + if "$ref" in content["schema"]: + # get the reference + root = openapi + reference = content["schema"]["$ref"].split("/")[1:] + for ref in reference: + root = root[ref] + # overwrite the content + interface["operation"]["requestBody"]["content"][content_type]["schema"] = root + + # parse body parameters + if "schema" in interface["operation"]["requestBody"]["content"][content_type]: + body_schema = interface["operation"]["requestBody"]["content"][content_type]["schema"] + required = body_schema.get("required", []) + properties = body_schema.get("properties", {}) + for name, property in properties.items(): + tool = ToolParameter( + name=name, + label=I18nObject(en_US=name, zh_Hans=name), + human_description=I18nObject( + en_US=property.get("description", ""), zh_Hans=property.get("description", "") + ), + type=ToolParameter.ToolParameterType.STRING, + required=name in required, + form=ToolParameter.ToolParameterForm.LLM, + llm_description=property.get("description", ""), + default=property.get("default", None), + placeholder=I18nObject( + en_US=property.get("description", ""), zh_Hans=property.get("description", "") + ), + ) + + # check if there is a type + typ = ApiBasedToolSchemaParser._get_tool_parameter_type(property) + if typ: + tool.type = typ + + parameters.append(tool) + + # check if parameters is duplicated + parameters_count = {} + for parameter in parameters: + if parameter.name not in parameters_count: + parameters_count[parameter.name] = 0 + parameters_count[parameter.name] += 1 + for name, count in parameters_count.items(): + if count > 1: + warning["duplicated_parameter"] = f"Parameter {name} is duplicated." + + # check if there is a operation id, use $path_$method as operation id if not + if "operationId" not in interface["operation"]: + # remove special characters like / to ensure the operation id is valid ^[a-zA-Z0-9_-]{1,64}$ + path = interface["path"] + if interface["path"].startswith("/"): + path = interface["path"][1:] + # remove special characters like / to ensure the operation id is valid ^[a-zA-Z0-9_-]{1,64}$ + path = re.sub(r"[^a-zA-Z0-9_-]", "", path) + if not path: + path = str(uuid.uuid4()) + + interface["operation"]["operationId"] = f"{path}_{interface['method']}" + + bundles.append( + ApiToolBundle( + server_url=server_url + interface["path"], + method=interface["method"], + summary=interface["operation"]["description"] + if "description" in interface["operation"] + else interface["operation"].get("summary", None), + operation_id=interface["operation"]["operationId"], + parameters=parameters, + author="", + icon=None, + openapi=interface["operation"], + ) + ) + + return bundles + + @staticmethod + def _get_tool_parameter_type(parameter: dict) -> Optional[ToolParameter.ToolParameterType]: + parameter = parameter or {} + typ: Optional[str] = None + if parameter.get("format") == "binary": + return ToolParameter.ToolParameterType.FILE + + if "type" in parameter: + typ = parameter["type"] + elif "schema" in parameter and "type" in parameter["schema"]: + typ = parameter["schema"]["type"] + + if typ in {"integer", "number"}: + return ToolParameter.ToolParameterType.NUMBER + elif typ == "boolean": + return ToolParameter.ToolParameterType.BOOLEAN + elif typ == "string": + return ToolParameter.ToolParameterType.STRING + else: + return None + + @staticmethod + def parse_openapi_yaml_to_tool_bundle( + yaml: str, extra_info: Optional[dict], warning: Optional[dict] + ) -> list[ApiToolBundle]: + """ + parse openapi yaml to tool bundle + + :param yaml: the yaml string + :return: the tool bundle + """ + warning = warning if warning is not None else {} + extra_info = extra_info if extra_info is not None else {} + + openapi: dict = safe_load(yaml) + if openapi is None: + raise ToolApiSchemaError("Invalid openapi yaml.") + return ApiBasedToolSchemaParser.parse_openapi_to_tool_bundle(openapi, extra_info=extra_info, warning=warning) + + @staticmethod + def parse_swagger_to_openapi(swagger: dict, extra_info: Optional[dict], warning: Optional[dict]) -> dict: + """ + parse swagger to openapi + + :param swagger: the swagger dict + :return: the openapi dict + """ + # convert swagger to openapi + info = swagger.get("info", {"title": "Swagger", "description": "Swagger", "version": "1.0.0"}) + + servers = swagger.get("servers", []) + + if len(servers) == 0: + raise ToolApiSchemaError("No server found in the swagger yaml.") + + openapi = { + "openapi": "3.0.0", + "info": { + "title": info.get("title", "Swagger"), + "description": info.get("description", "Swagger"), + "version": info.get("version", "1.0.0"), + }, + "servers": swagger["servers"], + "paths": {}, + "components": {"schemas": {}}, + } + + # check paths + if "paths" not in swagger or len(swagger["paths"]) == 0: + raise ToolApiSchemaError("No paths found in the swagger yaml.") + + # convert paths + for path, path_item in swagger["paths"].items(): + openapi["paths"][path] = {} + for method, operation in path_item.items(): + if "operationId" not in operation: + raise ToolApiSchemaError(f"No operationId found in operation {method} {path}.") + + if ("summary" not in operation or len(operation["summary"]) == 0) and ( + "description" not in operation or len(operation["description"]) == 0 + ): + if warning is not None: + warning["missing_summary"] = f"No summary or description found in operation {method} {path}." + + openapi["paths"][path][method] = { + "operationId": operation["operationId"], + "summary": operation.get("summary", ""), + "description": operation.get("description", ""), + "parameters": operation.get("parameters", []), + "responses": operation.get("responses", {}), + } + + if "requestBody" in operation: + openapi["paths"][path][method]["requestBody"] = operation["requestBody"] + + # convert definitions + for name, definition in swagger["definitions"].items(): + openapi["components"]["schemas"][name] = definition + + return openapi + + @staticmethod + def parse_openai_plugin_json_to_tool_bundle( + json: str, extra_info: Optional[dict], warning: Optional[dict] + ) -> list[ApiToolBundle]: + """ + parse openapi plugin yaml to tool bundle + + :param json: the json string + :return: the tool bundle + """ + warning = warning if warning is not None else {} + extra_info = extra_info if extra_info is not None else {} + + try: + openai_plugin = json_loads(json) + api = openai_plugin["api"] + api_url = api["url"] + api_type = api["type"] + except: + raise ToolProviderNotFoundError("Invalid openai plugin json.") + + if api_type != "openapi": + raise ToolNotSupportedError("Only openapi is supported now.") + + # get openapi yaml + response = get(api_url, headers={"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "}, timeout=5) + + if response.status_code != 200: + raise ToolProviderNotFoundError("cannot get openapi yaml from url.") + + return ApiBasedToolSchemaParser.parse_openapi_yaml_to_tool_bundle( + response.text, extra_info=extra_info, warning=warning + ) + + @staticmethod + def auto_parse_to_tool_bundle( + content: str, extra_info: Optional[dict] = None, warning: Optional[dict] = None + ) -> tuple[list[ApiToolBundle], str]: + """ + auto parse to tool bundle + + :param content: the content + :return: tools bundle, schema_type + """ + warning = warning if warning is not None else {} + extra_info = extra_info if extra_info is not None else {} + + content = content.strip() + loaded_content = None + json_error = None + yaml_error = None + + try: + loaded_content = json_loads(content) + except JSONDecodeError as e: + json_error = e + + if loaded_content is None: + try: + loaded_content = safe_load(content) + except YAMLError as e: + yaml_error = e + if loaded_content is None: + raise ToolApiSchemaError( + f"Invalid api schema, schema is neither json nor yaml. json error: {str(json_error)}," + f" yaml error: {str(yaml_error)}" + ) + + swagger_error = None + openapi_error = None + openapi_plugin_error = None + schema_type = None + + try: + openapi = ApiBasedToolSchemaParser.parse_openapi_to_tool_bundle( + loaded_content, extra_info=extra_info, warning=warning + ) + schema_type = ApiProviderSchemaType.OPENAPI.value + return openapi, schema_type + except ToolApiSchemaError as e: + openapi_error = e + + # openai parse error, fallback to swagger + try: + converted_swagger = ApiBasedToolSchemaParser.parse_swagger_to_openapi( + loaded_content, extra_info=extra_info, warning=warning + ) + schema_type = ApiProviderSchemaType.SWAGGER.value + return ApiBasedToolSchemaParser.parse_openapi_to_tool_bundle( + converted_swagger, extra_info=extra_info, warning=warning + ), schema_type + except ToolApiSchemaError as e: + swagger_error = e + + # swagger parse error, fallback to openai plugin + try: + openapi_plugin = ApiBasedToolSchemaParser.parse_openai_plugin_json_to_tool_bundle( + json_dumps(loaded_content), extra_info=extra_info, warning=warning + ) + return openapi_plugin, ApiProviderSchemaType.OPENAI_PLUGIN.value + except ToolNotSupportedError as e: + # maybe it's not plugin at all + openapi_plugin_error = e + + raise ToolApiSchemaError( + f"Invalid api schema, openapi error: {str(openapi_error)}, swagger error: {str(swagger_error)}," + f" openapi plugin error: {str(openapi_plugin_error)}" + ) diff --git a/api/core/tools/utils/text_processing_utils.py b/api/core/tools/utils/text_processing_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..105823f896788ef955afddb8af469ecf13f28e04 --- /dev/null +++ b/api/core/tools/utils/text_processing_utils.py @@ -0,0 +1,17 @@ +import re + + +def remove_leading_symbols(text: str) -> str: + """ + Remove leading punctuation or symbols from the given text. + + Args: + text (str): The input text to process. + + Returns: + str: The text with leading punctuation or symbols removed. + """ + # Match Unicode ranges for punctuation and symbols + # FIXME this pattern is confused quick fix for #11868 maybe refactor it later + pattern = r"^[\u2000-\u206F\u2E00-\u2E7F\u3000-\u303F!\"#$%&'()*+,./:;<=>?@^_`~]+" + return re.sub(pattern, "", text) diff --git a/api/core/tools/utils/uuid_utils.py b/api/core/tools/utils/uuid_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3046c08c89f0af29988f635d93ec01e3b46ea016 --- /dev/null +++ b/api/core/tools/utils/uuid_utils.py @@ -0,0 +1,9 @@ +import uuid + + +def is_valid_uuid(uuid_str: str) -> bool: + try: + uuid.UUID(uuid_str) + return True + except Exception: + return False diff --git a/api/core/tools/utils/web_reader_tool.py b/api/core/tools/utils/web_reader_tool.py new file mode 100644 index 0000000000000000000000000000000000000000..d42fd99fce5e801b04e703a74979e29d18279777 --- /dev/null +++ b/api/core/tools/utils/web_reader_tool.py @@ -0,0 +1,375 @@ +import hashlib +import json +import mimetypes +import os +import re +import site +import subprocess +import tempfile +import unicodedata +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Literal, Optional, cast +from urllib.parse import unquote + +import chardet +import cloudscraper # type: ignore +from bs4 import BeautifulSoup, CData, Comment, NavigableString # type: ignore +from regex import regex # type: ignore + +from core.helper import ssrf_proxy +from core.rag.extractor import extract_processor +from core.rag.extractor.extract_processor import ExtractProcessor + +FULL_TEMPLATE = """ +TITLE: {title} +AUTHORS: {authors} +PUBLISH DATE: {publish_date} +TOP_IMAGE_URL: {top_image} +TEXT: + +{text} +""" + + +def page_result(text: str, cursor: int, max_length: int) -> str: + """Page through `text` and return a substring of `max_length` characters starting from `cursor`.""" + return text[cursor : cursor + max_length] + + +def get_url(url: str, user_agent: Optional[str] = None) -> str: + """Fetch URL and return the contents as a string.""" + headers = { + "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko)" + " Chrome/91.0.4472.124 Safari/537.36" + } + if user_agent: + headers["User-Agent"] = user_agent + + main_content_type = None + supported_content_types = extract_processor.SUPPORT_URL_CONTENT_TYPES + ["text/html"] + response = ssrf_proxy.head(url, headers=headers, follow_redirects=True, timeout=(5, 10)) + + if response.status_code == 200: + # check content-type + content_type = response.headers.get("Content-Type") + if content_type: + main_content_type = response.headers.get("Content-Type").split(";")[0].strip() + else: + content_disposition = response.headers.get("Content-Disposition", "") + filename_match = re.search(r'filename="([^"]+)"', content_disposition) + if filename_match: + filename = unquote(filename_match.group(1)) + extension = re.search(r"\.(\w+)$", filename) + if extension: + main_content_type = mimetypes.guess_type(filename)[0] + + if main_content_type not in supported_content_types: + return "Unsupported content-type [{}] of URL.".format(main_content_type) + + if main_content_type in extract_processor.SUPPORT_URL_CONTENT_TYPES: + return cast(str, ExtractProcessor.load_from_url(url, return_text=True)) + + response = ssrf_proxy.get(url, headers=headers, follow_redirects=True, timeout=(120, 300)) + elif response.status_code == 403: + scraper = cloudscraper.create_scraper() + scraper.perform_request = ssrf_proxy.make_request + response = scraper.get(url, headers=headers, follow_redirects=True, timeout=(120, 300)) + + if response.status_code != 200: + return "URL returned status code {}.".format(response.status_code) + + # Detect encoding using chardet + detected_encoding = chardet.detect(response.content) + encoding = detected_encoding["encoding"] + if encoding: + try: + content = response.content.decode(encoding) + except (UnicodeDecodeError, TypeError): + content = response.text + else: + content = response.text + + a = extract_using_readabilipy(content) + + if not a["plain_text"] or not a["plain_text"].strip(): + return "" + + res = FULL_TEMPLATE.format( + title=a["title"], + authors=a["byline"], + publish_date=a["date"], + top_image="", + text=a["plain_text"] or "", + ) + + return res + + +def extract_using_readabilipy(html): + with tempfile.NamedTemporaryFile(delete=False, mode="w+") as f_html: + f_html.write(html) + f_html.close() + html_path = f_html.name + + # Call Mozilla's Readability.js Readability.parse() function via node, writing output to a temporary file + article_json_path = html_path + ".json" + jsdir = os.path.join(find_module_path("readabilipy"), "javascript") + with chdir(jsdir): + subprocess.check_call(["node", "ExtractArticle.js", "-i", html_path, "-o", article_json_path]) + + # Read output of call to Readability.parse() from JSON file and return as Python dictionary + input_json = json.loads(Path(article_json_path).read_text(encoding="utf-8")) + + # Deleting files after processing + os.unlink(article_json_path) + os.unlink(html_path) + + article_json: dict[str, Any] = { + "title": None, + "byline": None, + "date": None, + "content": None, + "plain_content": None, + "plain_text": None, + } + # Populate article fields from readability fields where present + if input_json: + if input_json.get("title"): + article_json["title"] = input_json["title"] + if input_json.get("byline"): + article_json["byline"] = input_json["byline"] + if input_json.get("date"): + article_json["date"] = input_json["date"] + if input_json.get("content"): + article_json["content"] = input_json["content"] + article_json["plain_content"] = plain_content(article_json["content"], False, False) + article_json["plain_text"] = extract_text_blocks_as_plain_text(article_json["plain_content"]) + if input_json.get("textContent"): + article_json["plain_text"] = input_json["textContent"] + article_json["plain_text"] = re.sub(r"\n\s*\n", "\n", article_json["plain_text"]) + + return article_json + + +def find_module_path(module_name): + for package_path in site.getsitepackages(): + potential_path = os.path.join(package_path, module_name) + if os.path.exists(potential_path): + return potential_path + + return None + + +@contextmanager +def chdir(path): + """Change directory in context and return to original on exit""" + # From https://stackoverflow.com/a/37996581, couldn't find a built-in + original_path = os.getcwd() + os.chdir(path) + try: + yield + finally: + os.chdir(original_path) + + +def extract_text_blocks_as_plain_text(paragraph_html): + # Load article as DOM + soup = BeautifulSoup(paragraph_html, "html.parser") + # Select all lists + list_elements = soup.find_all(["ul", "ol"]) + # Prefix text in all list items with "* " and make lists paragraphs + for list_element in list_elements: + plain_items = "".join( + list(filter(None, [plain_text_leaf_node(li)["text"] for li in list_element.find_all("li")])) + ) + list_element.string = plain_items + list_element.name = "p" + # Select all text blocks + text_blocks = [s.parent for s in soup.find_all(string=True)] + text_blocks = [plain_text_leaf_node(block) for block in text_blocks] + # Drop empty paragraphs + text_blocks = list(filter(lambda p: p["text"] is not None, text_blocks)) + return text_blocks + + +def plain_text_leaf_node(element): + # Extract all text, stripped of any child HTML elements and normalize it + plain_text = normalize_text(element.get_text()) + if plain_text != "" and element.name == "li": + plain_text = "* {}, ".format(plain_text) + if plain_text == "": + plain_text = None + if "data-node-index" in element.attrs: + plain = {"node_index": element["data-node-index"], "text": plain_text} + else: + plain = {"text": plain_text} + return plain + + +def plain_content(readability_content, content_digests, node_indexes): + # Load article as DOM + soup = BeautifulSoup(readability_content, "html.parser") + # Make all elements plain + elements = plain_elements(soup.contents, content_digests, node_indexes) + if node_indexes: + # Add node index attributes to nodes + elements = [add_node_indexes(element) for element in elements] + # Replace article contents with plain elements + soup.contents = elements + return str(soup) + + +def plain_elements(elements, content_digests, node_indexes): + # Get plain content versions of all elements + elements = [plain_element(element, content_digests, node_indexes) for element in elements] + if content_digests: + # Add content digest attribute to nodes + elements = [add_content_digest(element) for element in elements] + return elements + + +def plain_element(element, content_digests, node_indexes): + # For lists, we make each item plain text + if is_leaf(element): + # For leaf node elements, extract the text content, discarding any HTML tags + # 1. Get element contents as text + plain_text = element.get_text() + # 2. Normalize the extracted text string to a canonical representation + plain_text = normalize_text(plain_text) + # 3. Update element content to be plain text + element.string = plain_text + elif is_text(element): + if is_non_printing(element): + # The simplified HTML may have come from Readability.js so might + # have non-printing text (e.g. Comment or CData). In this case, we + # keep the structure, but ensure that the string is empty. + element = type(element)("") + else: + plain_text = element.string + plain_text = normalize_text(plain_text) + element = type(element)(plain_text) + else: + # If not a leaf node or leaf type call recursively on child nodes, replacing + element.contents = plain_elements(element.contents, content_digests, node_indexes) + return element + + +def add_node_indexes(element, node_index="0"): + # Can't add attributes to string types + if is_text(element): + return element + # Add index to current element + element["data-node-index"] = node_index + # Add index to child elements + for local_idx, child in enumerate([c for c in element.contents if not is_text(c)], start=1): + # Can't add attributes to leaf string types + child_index = "{stem}.{local}".format(stem=node_index, local=local_idx) + add_node_indexes(child, node_index=child_index) + return element + + +def normalize_text(text): + """Normalize unicode and whitespace.""" + # Normalize unicode first to try and standardize whitespace characters as much as possible before normalizing them + text = strip_control_characters(text) + text = normalize_unicode(text) + text = normalize_whitespace(text) + return text + + +def strip_control_characters(text): + """Strip out unicode control characters which might break the parsing.""" + # Unicode control characters + # [Cc]: Other, Control [includes new lines] + # [Cf]: Other, Format + # [Cn]: Other, Not Assigned + # [Co]: Other, Private Use + # [Cs]: Other, Surrogate + control_chars = {"Cc", "Cf", "Cn", "Co", "Cs"} + retained_chars = ["\t", "\n", "\r", "\f"] + + # Remove non-printing control characters + return "".join( + [ + "" if (unicodedata.category(char) in control_chars) and (char not in retained_chars) else char + for char in text + ] + ) + + +def normalize_unicode(text): + """Normalize unicode such that things that are visually equivalent map to the same unicode string where possible.""" + normal_form: Literal["NFC", "NFD", "NFKC", "NFKD"] = "NFKC" + text = unicodedata.normalize(normal_form, text) + return text + + +def normalize_whitespace(text): + """Replace runs of whitespace characters with a single space as this is what happens when HTML text is displayed.""" + text = regex.sub(r"\s+", " ", text) + # Remove leading and trailing whitespace + text = text.strip() + return text + + +def is_leaf(element): + return element.name in {"p", "li"} + + +def is_text(element): + return isinstance(element, NavigableString) + + +def is_non_printing(element): + return any(isinstance(element, _e) for _e in [Comment, CData]) + + +def add_content_digest(element): + if not is_text(element): + element["data-content-digest"] = content_digest(element) + return element + + +def content_digest(element): + digest: Any + if is_text(element): + # Hash + trimmed_string = element.string.strip() + if trimmed_string == "": + digest = "" + else: + digest = hashlib.sha256(trimmed_string.encode("utf-8")).hexdigest() + else: + contents = element.contents + num_contents = len(contents) + if num_contents == 0: + # No hash when no child elements exist + digest = "" + elif num_contents == 1: + # If single child, use digest of child + digest = content_digest(contents[0]) + else: + # Build content digest from the "non-empty" digests of child nodes + digest = hashlib.sha256() + child_digests = list(filter(lambda x: x != "", [content_digest(content) for content in contents])) + for child in child_digests: + digest.update(child.encode("utf-8")) + digest = digest.hexdigest() + return digest + + +def get_image_upload_file_ids(content): + pattern = r"!\[image\]\((http?://.*?(file-preview|image-preview))\)" + matches = re.findall(pattern, content) + image_upload_file_ids = [] + for match in matches: + if match[1] == "file-preview": + content_pattern = r"files/([^/]+)/file-preview" + else: + content_pattern = r"files/([^/]+)/image-preview" + content_match = re.search(content_pattern, match[0]) + if content_match: + image_upload_file_id = content_match.group(1) + image_upload_file_ids.append(image_upload_file_id) + return image_upload_file_ids diff --git a/api/core/tools/utils/workflow_configuration_sync.py b/api/core/tools/utils/workflow_configuration_sync.py new file mode 100644 index 0000000000000000000000000000000000000000..08a112cfdb2b917c9575f549c1bc9e0c2f9a8c04 --- /dev/null +++ b/api/core/tools/utils/workflow_configuration_sync.py @@ -0,0 +1,45 @@ +from collections.abc import Mapping, Sequence +from typing import Any + +from core.app.app_config.entities import VariableEntity +from core.tools.entities.tool_entities import WorkflowToolParameterConfiguration + + +class WorkflowToolConfigurationUtils: + @classmethod + def check_parameter_configurations(cls, configurations: list[Mapping[str, Any]]): + for configuration in configurations: + WorkflowToolParameterConfiguration.model_validate(configuration) + + @classmethod + def get_workflow_graph_variables(cls, graph: Mapping[str, Any]) -> Sequence[VariableEntity]: + """ + get workflow graph variables + """ + nodes = graph.get("nodes", []) + start_node = next(filter(lambda x: x.get("data", {}).get("type") == "start", nodes), None) + + if not start_node: + return [] + + return [VariableEntity.model_validate(variable) for variable in start_node.get("data", {}).get("variables", [])] + + @classmethod + def check_is_synced( + cls, variables: list[VariableEntity], tool_configurations: list[WorkflowToolParameterConfiguration] + ) -> bool: + """ + check is synced + + raise ValueError if not synced + """ + variable_names = [variable.variable for variable in variables] + + if len(tool_configurations) != len(variables): + raise ValueError("parameter configuration mismatch, please republish the tool to update") + + for parameter in tool_configurations: + if parameter.name not in variable_names: + raise ValueError("parameter configuration mismatch, please republish the tool to update") + + return True diff --git a/api/core/tools/utils/yaml_utils.py b/api/core/tools/utils/yaml_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ee7ca11e056625bc4e828e3b1794d4d71aa38c50 --- /dev/null +++ b/api/core/tools/utils/yaml_utils.py @@ -0,0 +1,35 @@ +import logging +from pathlib import Path +from typing import Any + +import yaml # type: ignore +from yaml import YAMLError + +logger = logging.getLogger(__name__) + + +def load_yaml_file(file_path: str, ignore_error: bool = True, default_value: Any = {}) -> Any: + """ + Safe loading a YAML file + :param file_path: the path of the YAML file + :param ignore_error: + if True, return default_value if error occurs and the error will be logged in debug level + if False, raise error if error occurs + :param default_value: the value returned when errors ignored + :return: an object of the YAML content + """ + if not file_path or not Path(file_path).exists(): + if ignore_error: + return default_value + else: + raise FileNotFoundError(f"File not found: {file_path}") + + with open(file_path, encoding="utf-8") as yaml_file: + try: + yaml_content = yaml.safe_load(yaml_file) + return yaml_content or default_value + except Exception as e: + if ignore_error: + return default_value + else: + raise YAMLError(f"Failed to load YAML file {file_path}: {e}") from e diff --git a/api/core/variables/__init__.py b/api/core/variables/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7a1cbf99407ea8c4ad04af59a378fec6fb74cc70 --- /dev/null +++ b/api/core/variables/__init__.py @@ -0,0 +1,65 @@ +from .segment_group import SegmentGroup +from .segments import ( + ArrayAnySegment, + ArrayFileSegment, + ArrayNumberSegment, + ArrayObjectSegment, + ArraySegment, + ArrayStringSegment, + FileSegment, + FloatSegment, + IntegerSegment, + NoneSegment, + ObjectSegment, + Segment, + StringSegment, +) +from .types import SegmentType +from .variables import ( + ArrayAnyVariable, + ArrayFileVariable, + ArrayNumberVariable, + ArrayObjectVariable, + ArrayStringVariable, + ArrayVariable, + FileVariable, + FloatVariable, + IntegerVariable, + NoneVariable, + ObjectVariable, + SecretVariable, + StringVariable, + Variable, +) + +__all__ = [ + "ArrayAnySegment", + "ArrayAnyVariable", + "ArrayFileSegment", + "ArrayFileVariable", + "ArrayNumberSegment", + "ArrayNumberVariable", + "ArrayObjectSegment", + "ArrayObjectVariable", + "ArraySegment", + "ArrayStringSegment", + "ArrayStringVariable", + "ArrayVariable", + "FileSegment", + "FileVariable", + "FloatSegment", + "FloatVariable", + "IntegerSegment", + "IntegerVariable", + "NoneSegment", + "NoneVariable", + "ObjectSegment", + "ObjectVariable", + "SecretVariable", + "Segment", + "SegmentGroup", + "SegmentType", + "StringSegment", + "StringVariable", + "Variable", +] diff --git a/api/core/variables/exc.py b/api/core/variables/exc.py new file mode 100644 index 0000000000000000000000000000000000000000..5cf67c3baccacc610270bf3ba92a0db73fbd9b5a --- /dev/null +++ b/api/core/variables/exc.py @@ -0,0 +1,2 @@ +class VariableError(ValueError): + pass diff --git a/api/core/variables/segment_group.py b/api/core/variables/segment_group.py new file mode 100644 index 0000000000000000000000000000000000000000..b363255b2cae9e8bf47e4500dff8a662b5be4a71 --- /dev/null +++ b/api/core/variables/segment_group.py @@ -0,0 +1,22 @@ +from .segments import Segment +from .types import SegmentType + + +class SegmentGroup(Segment): + value_type: SegmentType = SegmentType.GROUP + value: list[Segment] + + @property + def text(self): + return "".join([segment.text for segment in self.value]) + + @property + def log(self): + return "".join([segment.log for segment in self.value]) + + @property + def markdown(self): + return "".join([segment.markdown for segment in self.value]) + + def to_object(self): + return [segment.to_object() for segment in self.value] diff --git a/api/core/variables/segments.py b/api/core/variables/segments.py new file mode 100644 index 0000000000000000000000000000000000000000..a9f5651692bb4a16b7a8f1d639804628bd310df5 --- /dev/null +++ b/api/core/variables/segments.py @@ -0,0 +1,169 @@ +import json +import sys +from collections.abc import Mapping, Sequence +from typing import Any + +from pydantic import BaseModel, ConfigDict, field_validator + +from core.file import File + +from .types import SegmentType + + +class Segment(BaseModel): + model_config = ConfigDict(frozen=True) + + value_type: SegmentType + value: Any + + @field_validator("value_type") + @classmethod + def validate_value_type(cls, value): + """ + This validator checks if the provided value is equal to the default value of the 'value_type' field. + If the value is different, a ValueError is raised. + """ + if value != cls.model_fields["value_type"].default: + raise ValueError("Cannot modify 'value_type'") + return value + + @property + def text(self) -> str: + return str(self.value) + + @property + def log(self) -> str: + return str(self.value) + + @property + def markdown(self) -> str: + return str(self.value) + + @property + def size(self) -> int: + """ + Return the size of the value in bytes. + """ + return sys.getsizeof(self.value) + + def to_object(self) -> Any: + return self.value + + +class NoneSegment(Segment): + value_type: SegmentType = SegmentType.NONE + value: None = None + + @property + def text(self) -> str: + return "" + + @property + def log(self) -> str: + return "" + + @property + def markdown(self) -> str: + return "" + + +class StringSegment(Segment): + value_type: SegmentType = SegmentType.STRING + value: str + + +class FloatSegment(Segment): + value_type: SegmentType = SegmentType.NUMBER + value: float + + +class IntegerSegment(Segment): + value_type: SegmentType = SegmentType.NUMBER + value: int + + +class ObjectSegment(Segment): + value_type: SegmentType = SegmentType.OBJECT + value: Mapping[str, Any] + + @property + def text(self) -> str: + return json.dumps(self.model_dump()["value"], ensure_ascii=False) + + @property + def log(self) -> str: + return json.dumps(self.model_dump()["value"], ensure_ascii=False, indent=2) + + @property + def markdown(self) -> str: + return json.dumps(self.model_dump()["value"], ensure_ascii=False, indent=2) + + +class ArraySegment(Segment): + @property + def markdown(self) -> str: + items = [] + for item in self.value: + items.append(str(item)) + return "\n".join(items) + + +class FileSegment(Segment): + value_type: SegmentType = SegmentType.FILE + value: File + + @property + def markdown(self) -> str: + return self.value.markdown + + @property + def log(self) -> str: + return "" + + @property + def text(self) -> str: + return "" + + +class ArrayAnySegment(ArraySegment): + value_type: SegmentType = SegmentType.ARRAY_ANY + value: Sequence[Any] + + +class ArrayStringSegment(ArraySegment): + value_type: SegmentType = SegmentType.ARRAY_STRING + value: Sequence[str] + + @property + def text(self) -> str: + return json.dumps(self.value) + + +class ArrayNumberSegment(ArraySegment): + value_type: SegmentType = SegmentType.ARRAY_NUMBER + value: Sequence[float | int] + + +class ArrayObjectSegment(ArraySegment): + value_type: SegmentType = SegmentType.ARRAY_OBJECT + value: Sequence[Mapping[str, Any]] + + +class ArrayFileSegment(ArraySegment): + value_type: SegmentType = SegmentType.ARRAY_FILE + value: Sequence[File] + + @property + def markdown(self) -> str: + items = [] + for item in self.value: + items.append(item.markdown) + return "\n".join(items) + + @property + def log(self) -> str: + return "" + + @property + def text(self) -> str: + return "" diff --git a/api/core/variables/types.py b/api/core/variables/types.py new file mode 100644 index 0000000000000000000000000000000000000000..4387e9693eb072ff55938d99e94f933e3c95cb6d --- /dev/null +++ b/api/core/variables/types.py @@ -0,0 +1,20 @@ +from enum import StrEnum + + +class SegmentType(StrEnum): + NUMBER = "number" + STRING = "string" + OBJECT = "object" + SECRET = "secret" + + FILE = "file" + + ARRAY_ANY = "array[any]" + ARRAY_STRING = "array[string]" + ARRAY_NUMBER = "array[number]" + ARRAY_OBJECT = "array[object]" + ARRAY_FILE = "array[file]" + + NONE = "none" + + GROUP = "group" diff --git a/api/core/variables/variables.py b/api/core/variables/variables.py new file mode 100644 index 0000000000000000000000000000000000000000..c32815b24d02ed60e90fe35ebe7110352ec3a0db --- /dev/null +++ b/api/core/variables/variables.py @@ -0,0 +1,95 @@ +from collections.abc import Sequence +from typing import cast +from uuid import uuid4 + +from pydantic import Field + +from core.helper import encrypter + +from .segments import ( + ArrayAnySegment, + ArrayFileSegment, + ArrayNumberSegment, + ArrayObjectSegment, + ArraySegment, + ArrayStringSegment, + FileSegment, + FloatSegment, + IntegerSegment, + NoneSegment, + ObjectSegment, + Segment, + StringSegment, +) +from .types import SegmentType + + +class Variable(Segment): + """ + A variable is a segment that has a name. + """ + + id: str = Field( + default=lambda _: str(uuid4()), + description="Unique identity for variable.", + ) + name: str + description: str = Field(default="", description="Description of the variable.") + selector: Sequence[str] = Field(default_factory=list) + + +class StringVariable(StringSegment, Variable): + pass + + +class FloatVariable(FloatSegment, Variable): + pass + + +class IntegerVariable(IntegerSegment, Variable): + pass + + +class ObjectVariable(ObjectSegment, Variable): + pass + + +class ArrayVariable(ArraySegment, Variable): + pass + + +class ArrayAnyVariable(ArrayAnySegment, ArrayVariable): + pass + + +class ArrayStringVariable(ArrayStringSegment, ArrayVariable): + pass + + +class ArrayNumberVariable(ArrayNumberSegment, ArrayVariable): + pass + + +class ArrayObjectVariable(ArrayObjectSegment, ArrayVariable): + pass + + +class SecretVariable(StringVariable): + value_type: SegmentType = SegmentType.SECRET + + @property + def log(self) -> str: + return cast(str, encrypter.obfuscated_token(self.value)) + + +class NoneVariable(NoneSegment, Variable): + value_type: SegmentType = SegmentType.NONE + value: None = None + + +class FileVariable(FileSegment, Variable): + pass + + +class ArrayFileVariable(ArrayFileSegment, ArrayVariable): + pass diff --git a/api/core/workflow/__init__.py b/api/core/workflow/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/workflow/callbacks/__init__.py b/api/core/workflow/callbacks/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fba86c1e2eb2ea19a4cf176c3610838fa63c1d35 --- /dev/null +++ b/api/core/workflow/callbacks/__init__.py @@ -0,0 +1,7 @@ +from .base_workflow_callback import WorkflowCallback +from .workflow_logging_callback import WorkflowLoggingCallback + +__all__ = [ + "WorkflowCallback", + "WorkflowLoggingCallback", +] diff --git a/api/core/workflow/callbacks/base_workflow_callback.py b/api/core/workflow/callbacks/base_workflow_callback.py new file mode 100644 index 0000000000000000000000000000000000000000..83086d1afc9018273be75be65d132e707dac3c43 --- /dev/null +++ b/api/core/workflow/callbacks/base_workflow_callback.py @@ -0,0 +1,12 @@ +from abc import ABC, abstractmethod + +from core.workflow.graph_engine.entities.event import GraphEngineEvent + + +class WorkflowCallback(ABC): + @abstractmethod + def on_event(self, event: GraphEngineEvent) -> None: + """ + Published event + """ + raise NotImplementedError diff --git a/api/core/workflow/callbacks/workflow_logging_callback.py b/api/core/workflow/callbacks/workflow_logging_callback.py new file mode 100644 index 0000000000000000000000000000000000000000..b9c6b35ad3476ada48dad87a6ad0600f7e093353 --- /dev/null +++ b/api/core/workflow/callbacks/workflow_logging_callback.py @@ -0,0 +1,224 @@ +from typing import Optional + +from core.model_runtime.utils.encoders import jsonable_encoder +from core.workflow.graph_engine.entities.event import ( + GraphEngineEvent, + GraphRunFailedEvent, + GraphRunPartialSucceededEvent, + GraphRunStartedEvent, + GraphRunSucceededEvent, + IterationRunFailedEvent, + IterationRunNextEvent, + IterationRunStartedEvent, + IterationRunSucceededEvent, + NodeRunFailedEvent, + NodeRunStartedEvent, + NodeRunStreamChunkEvent, + NodeRunSucceededEvent, + ParallelBranchRunFailedEvent, + ParallelBranchRunStartedEvent, + ParallelBranchRunSucceededEvent, +) + +from .base_workflow_callback import WorkflowCallback + +_TEXT_COLOR_MAPPING = { + "blue": "36;1", + "yellow": "33;1", + "pink": "38;5;200", + "green": "32;1", + "red": "31;1", +} + + +class WorkflowLoggingCallback(WorkflowCallback): + def __init__(self) -> None: + self.current_node_id: Optional[str] = None + + def on_event(self, event: GraphEngineEvent) -> None: + if isinstance(event, GraphRunStartedEvent): + self.print_text("\n[GraphRunStartedEvent]", color="pink") + elif isinstance(event, GraphRunSucceededEvent): + self.print_text("\n[GraphRunSucceededEvent]", color="green") + elif isinstance(event, GraphRunPartialSucceededEvent): + self.print_text("\n[GraphRunPartialSucceededEvent]", color="pink") + elif isinstance(event, GraphRunFailedEvent): + self.print_text(f"\n[GraphRunFailedEvent] reason: {event.error}", color="red") + elif isinstance(event, NodeRunStartedEvent): + self.on_workflow_node_execute_started(event=event) + elif isinstance(event, NodeRunSucceededEvent): + self.on_workflow_node_execute_succeeded(event=event) + elif isinstance(event, NodeRunFailedEvent): + self.on_workflow_node_execute_failed(event=event) + elif isinstance(event, NodeRunStreamChunkEvent): + self.on_node_text_chunk(event=event) + elif isinstance(event, ParallelBranchRunStartedEvent): + self.on_workflow_parallel_started(event=event) + elif isinstance(event, ParallelBranchRunSucceededEvent | ParallelBranchRunFailedEvent): + self.on_workflow_parallel_completed(event=event) + elif isinstance(event, IterationRunStartedEvent): + self.on_workflow_iteration_started(event=event) + elif isinstance(event, IterationRunNextEvent): + self.on_workflow_iteration_next(event=event) + elif isinstance(event, IterationRunSucceededEvent | IterationRunFailedEvent): + self.on_workflow_iteration_completed(event=event) + else: + self.print_text(f"\n[{event.__class__.__name__}]", color="blue") + + def on_workflow_node_execute_started(self, event: NodeRunStartedEvent) -> None: + """ + Workflow node execute started + """ + self.print_text("\n[NodeRunStartedEvent]", color="yellow") + self.print_text(f"Node ID: {event.node_id}", color="yellow") + self.print_text(f"Node Title: {event.node_data.title}", color="yellow") + self.print_text(f"Type: {event.node_type.value}", color="yellow") + + def on_workflow_node_execute_succeeded(self, event: NodeRunSucceededEvent) -> None: + """ + Workflow node execute succeeded + """ + route_node_state = event.route_node_state + + self.print_text("\n[NodeRunSucceededEvent]", color="green") + self.print_text(f"Node ID: {event.node_id}", color="green") + self.print_text(f"Node Title: {event.node_data.title}", color="green") + self.print_text(f"Type: {event.node_type.value}", color="green") + + if route_node_state.node_run_result: + node_run_result = route_node_state.node_run_result + self.print_text( + f"Inputs: {jsonable_encoder(node_run_result.inputs) if node_run_result.inputs else ''}", + color="green", + ) + self.print_text( + f"Process Data: " + f"{jsonable_encoder(node_run_result.process_data) if node_run_result.process_data else ''}", + color="green", + ) + self.print_text( + f"Outputs: {jsonable_encoder(node_run_result.outputs) if node_run_result.outputs else ''}", + color="green", + ) + self.print_text( + f"Metadata: {jsonable_encoder(node_run_result.metadata) if node_run_result.metadata else ''}", + color="green", + ) + + def on_workflow_node_execute_failed(self, event: NodeRunFailedEvent) -> None: + """ + Workflow node execute failed + """ + route_node_state = event.route_node_state + + self.print_text("\n[NodeRunFailedEvent]", color="red") + self.print_text(f"Node ID: {event.node_id}", color="red") + self.print_text(f"Node Title: {event.node_data.title}", color="red") + self.print_text(f"Type: {event.node_type.value}", color="red") + + if route_node_state.node_run_result: + node_run_result = route_node_state.node_run_result + self.print_text(f"Error: {node_run_result.error}", color="red") + self.print_text( + f"Inputs: {jsonable_encoder(node_run_result.inputs) if node_run_result.inputs else ''}", + color="red", + ) + self.print_text( + f"Process Data: " + f"{jsonable_encoder(node_run_result.process_data) if node_run_result.process_data else ''}", + color="red", + ) + self.print_text( + f"Outputs: {jsonable_encoder(node_run_result.outputs) if node_run_result.outputs else ''}", + color="red", + ) + + def on_node_text_chunk(self, event: NodeRunStreamChunkEvent) -> None: + """ + Publish text chunk + """ + route_node_state = event.route_node_state + if not self.current_node_id or self.current_node_id != route_node_state.node_id: + self.current_node_id = route_node_state.node_id + self.print_text("\n[NodeRunStreamChunkEvent]") + self.print_text(f"Node ID: {route_node_state.node_id}") + + node_run_result = route_node_state.node_run_result + if node_run_result: + self.print_text( + f"Metadata: {jsonable_encoder(node_run_result.metadata) if node_run_result.metadata else ''}" + ) + + self.print_text(event.chunk_content, color="pink", end="") + + def on_workflow_parallel_started(self, event: ParallelBranchRunStartedEvent) -> None: + """ + Publish parallel started + """ + self.print_text("\n[ParallelBranchRunStartedEvent]", color="blue") + self.print_text(f"Parallel ID: {event.parallel_id}", color="blue") + self.print_text(f"Branch ID: {event.parallel_start_node_id}", color="blue") + if event.in_iteration_id: + self.print_text(f"Iteration ID: {event.in_iteration_id}", color="blue") + + def on_workflow_parallel_completed( + self, event: ParallelBranchRunSucceededEvent | ParallelBranchRunFailedEvent + ) -> None: + """ + Publish parallel completed + """ + if isinstance(event, ParallelBranchRunSucceededEvent): + color = "blue" + elif isinstance(event, ParallelBranchRunFailedEvent): + color = "red" + + self.print_text( + "\n[ParallelBranchRunSucceededEvent]" + if isinstance(event, ParallelBranchRunSucceededEvent) + else "\n[ParallelBranchRunFailedEvent]", + color=color, + ) + self.print_text(f"Parallel ID: {event.parallel_id}", color=color) + self.print_text(f"Branch ID: {event.parallel_start_node_id}", color=color) + if event.in_iteration_id: + self.print_text(f"Iteration ID: {event.in_iteration_id}", color=color) + + if isinstance(event, ParallelBranchRunFailedEvent): + self.print_text(f"Error: {event.error}", color=color) + + def on_workflow_iteration_started(self, event: IterationRunStartedEvent) -> None: + """ + Publish iteration started + """ + self.print_text("\n[IterationRunStartedEvent]", color="blue") + self.print_text(f"Iteration Node ID: {event.iteration_id}", color="blue") + + def on_workflow_iteration_next(self, event: IterationRunNextEvent) -> None: + """ + Publish iteration next + """ + self.print_text("\n[IterationRunNextEvent]", color="blue") + self.print_text(f"Iteration Node ID: {event.iteration_id}", color="blue") + self.print_text(f"Iteration Index: {event.index}", color="blue") + + def on_workflow_iteration_completed(self, event: IterationRunSucceededEvent | IterationRunFailedEvent) -> None: + """ + Publish iteration completed + """ + self.print_text( + "\n[IterationRunSucceededEvent]" + if isinstance(event, IterationRunSucceededEvent) + else "\n[IterationRunFailedEvent]", + color="blue", + ) + self.print_text(f"Node ID: {event.iteration_id}", color="blue") + + def print_text(self, text: str, color: Optional[str] = None, end: str = "\n") -> None: + """Print text with highlighting and no end characters.""" + text_to_print = self._get_colored_text(text, color) if color else text + print(f"{text_to_print}", end=end) + + def _get_colored_text(self, text: str, color: str) -> str: + """Get colored text.""" + color_str = _TEXT_COLOR_MAPPING[color] + return f"\u001b[{color_str}m\033[1;3m{text}\u001b[0m" diff --git a/api/core/workflow/constants.py b/api/core/workflow/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..e3fe17c2845837c6a4f79f660b5b5be500c607ac --- /dev/null +++ b/api/core/workflow/constants.py @@ -0,0 +1,3 @@ +SYSTEM_VARIABLE_NODE_ID = "sys" +ENVIRONMENT_VARIABLE_NODE_ID = "env" +CONVERSATION_VARIABLE_NODE_ID = "conversation" diff --git a/api/core/workflow/entities/__init__.py b/api/core/workflow/entities/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/workflow/entities/node_entities.py b/api/core/workflow/entities/node_entities.py new file mode 100644 index 0000000000000000000000000000000000000000..ae5f117bf9b121921bd722f66a2cb392911723e4 --- /dev/null +++ b/api/core/workflow/entities/node_entities.py @@ -0,0 +1,50 @@ +from collections.abc import Mapping +from enum import StrEnum +from typing import Any, Optional + +from pydantic import BaseModel + +from core.model_runtime.entities.llm_entities import LLMUsage +from models.workflow import WorkflowNodeExecutionStatus + + +class NodeRunMetadataKey(StrEnum): + """ + Node Run Metadata Key. + """ + + TOTAL_TOKENS = "total_tokens" + TOTAL_PRICE = "total_price" + CURRENCY = "currency" + TOOL_INFO = "tool_info" + ITERATION_ID = "iteration_id" + ITERATION_INDEX = "iteration_index" + PARALLEL_ID = "parallel_id" + PARALLEL_START_NODE_ID = "parallel_start_node_id" + PARENT_PARALLEL_ID = "parent_parallel_id" + PARENT_PARALLEL_START_NODE_ID = "parent_parallel_start_node_id" + PARALLEL_MODE_RUN_ID = "parallel_mode_run_id" + ITERATION_DURATION_MAP = "iteration_duration_map" # single iteration duration if iteration node runs + ERROR_STRATEGY = "error_strategy" # node in continue on error mode return the field + + +class NodeRunResult(BaseModel): + """ + Node Run Result. + """ + + status: WorkflowNodeExecutionStatus = WorkflowNodeExecutionStatus.RUNNING + + inputs: Optional[Mapping[str, Any]] = None # node inputs + process_data: Optional[Mapping[str, Any]] = None # process data + outputs: Optional[Mapping[str, Any]] = None # node outputs + metadata: Optional[Mapping[NodeRunMetadataKey, Any]] = None # node metadata + llm_usage: Optional[LLMUsage] = None # llm usage + + edge_source_handle: Optional[str] = None # source handle id of node with multiple branches + + error: Optional[str] = None # error message if status is failed + error_type: Optional[str] = None # error type if status is failed + + # single step node run retry + retry_index: int = 0 diff --git a/api/core/workflow/entities/variable_entities.py b/api/core/workflow/entities/variable_entities.py new file mode 100644 index 0000000000000000000000000000000000000000..8f4c2d797552ca7a3eb6e5aa5fc1233cd11bf978 --- /dev/null +++ b/api/core/workflow/entities/variable_entities.py @@ -0,0 +1,12 @@ +from collections.abc import Sequence + +from pydantic import BaseModel + + +class VariableSelector(BaseModel): + """ + Variable Selector. + """ + + variable: str + value_selector: Sequence[str] diff --git a/api/core/workflow/entities/variable_pool.py b/api/core/workflow/entities/variable_pool.py new file mode 100644 index 0000000000000000000000000000000000000000..844b46f3528d4258e7b7ad825afdcce9c35fb7aa --- /dev/null +++ b/api/core/workflow/entities/variable_pool.py @@ -0,0 +1,174 @@ +import re +from collections import defaultdict +from collections.abc import Mapping, Sequence +from typing import Any, Union + +from pydantic import BaseModel, Field + +from core.file import File, FileAttribute, file_manager +from core.variables import Segment, SegmentGroup, Variable +from core.variables.segments import FileSegment +from factories import variable_factory + +from ..constants import CONVERSATION_VARIABLE_NODE_ID, ENVIRONMENT_VARIABLE_NODE_ID, SYSTEM_VARIABLE_NODE_ID +from ..enums import SystemVariableKey + +VariableValue = Union[str, int, float, dict, list, File] + + +VARIABLE_PATTERN = re.compile(r"\{\{#([a-zA-Z0-9_]{1,50}(?:\.[a-zA-Z_][a-zA-Z0-9_]{0,29}){1,10})#\}\}") + + +class VariablePool(BaseModel): + # Variable dictionary is a dictionary for looking up variables by their selector. + # The first element of the selector is the node id, it's the first-level key in the dictionary. + # Other elements of the selector are the keys in the second-level dictionary. To get the key, we hash the + # elements of the selector except the first one. + variable_dictionary: dict[str, dict[int, Segment]] = Field( + description="Variables mapping", + default=defaultdict(dict), + ) + # TODO: This user inputs is not used for pool. + user_inputs: Mapping[str, Any] = Field( + description="User inputs", + ) + system_variables: Mapping[SystemVariableKey, Any] = Field( + description="System variables", + ) + environment_variables: Sequence[Variable] = Field( + description="Environment variables.", + default_factory=list, + ) + conversation_variables: Sequence[Variable] = Field( + description="Conversation variables.", + default_factory=list, + ) + + def __init__( + self, + *, + system_variables: Mapping[SystemVariableKey, Any] | None = None, + user_inputs: Mapping[str, Any] | None = None, + environment_variables: Sequence[Variable] | None = None, + conversation_variables: Sequence[Variable] | None = None, + **kwargs, + ): + environment_variables = environment_variables or [] + conversation_variables = conversation_variables or [] + user_inputs = user_inputs or {} + system_variables = system_variables or {} + + super().__init__( + system_variables=system_variables, + user_inputs=user_inputs, + environment_variables=environment_variables, + conversation_variables=conversation_variables, + **kwargs, + ) + + for key, value in self.system_variables.items(): + self.add((SYSTEM_VARIABLE_NODE_ID, key.value), value) + # Add environment variables to the variable pool + for var in self.environment_variables: + self.add((ENVIRONMENT_VARIABLE_NODE_ID, var.name), var) + # Add conversation variables to the variable pool + for var in self.conversation_variables: + self.add((CONVERSATION_VARIABLE_NODE_ID, var.name), var) + + def add(self, selector: Sequence[str], value: Any, /) -> None: + """ + Adds a variable to the variable pool. + + NOTE: You should not add a non-Segment value to the variable pool + even if it is allowed now. + + Args: + selector (Sequence[str]): The selector for the variable. + value (VariableValue): The value of the variable. + + Raises: + ValueError: If the selector is invalid. + + Returns: + None + """ + if len(selector) < 2: + raise ValueError("Invalid selector") + + if isinstance(value, Variable): + variable = value + if isinstance(value, Segment): + variable = variable_factory.segment_to_variable(segment=value, selector=selector) + else: + segment = variable_factory.build_segment(value) + variable = variable_factory.segment_to_variable(segment=segment, selector=selector) + + hash_key = hash(tuple(selector[1:])) + self.variable_dictionary[selector[0]][hash_key] = variable + + def get(self, selector: Sequence[str], /) -> Segment | None: + """ + Retrieves the value from the variable pool based on the given selector. + + Args: + selector (Sequence[str]): The selector used to identify the variable. + + Returns: + Any: The value associated with the given selector. + + Raises: + ValueError: If the selector is invalid. + """ + if len(selector) < 2: + return None + + hash_key = hash(tuple(selector[1:])) + value = self.variable_dictionary[selector[0]].get(hash_key) + + if value is None: + selector, attr = selector[:-1], selector[-1] + # Python support `attr in FileAttribute` after 3.12 + if attr not in {item.value for item in FileAttribute}: + return None + value = self.get(selector) + if not isinstance(value, FileSegment): + return None + attr = FileAttribute(attr) + attr_value = file_manager.get_attr(file=value.value, attr=attr) + return variable_factory.build_segment(attr_value) + + return value + + def remove(self, selector: Sequence[str], /): + """ + Remove variables from the variable pool based on the given selector. + + Args: + selector (Sequence[str]): A sequence of strings representing the selector. + + Returns: + None + """ + if not selector: + return + if len(selector) == 1: + self.variable_dictionary[selector[0]] = {} + return + hash_key = hash(tuple(selector[1:])) + self.variable_dictionary[selector[0]].pop(hash_key, None) + + def convert_template(self, template: str, /): + parts = VARIABLE_PATTERN.split(template) + segments = [] + for part in filter(lambda x: x, parts): + if "." in part and (variable := self.get(part.split("."))): + segments.append(variable) + else: + segments.append(variable_factory.build_segment(part)) + return SegmentGroup(value=segments) + + def get_file(self, selector: Sequence[str], /) -> FileSegment | None: + segment = self.get(selector) + if isinstance(segment, FileSegment): + return segment + return None diff --git a/api/core/workflow/entities/workflow_entities.py b/api/core/workflow/entities/workflow_entities.py new file mode 100644 index 0000000000000000000000000000000000000000..da56af1407d94fe902baee2e142cb38b195ad127 --- /dev/null +++ b/api/core/workflow/entities/workflow_entities.py @@ -0,0 +1,76 @@ +from typing import Optional + +from pydantic import BaseModel + +from core.app.entities.app_invoke_entities import InvokeFrom +from core.workflow.nodes.base import BaseIterationState, BaseNode +from models.enums import UserFrom +from models.workflow import Workflow, WorkflowType + +from .node_entities import NodeRunResult +from .variable_pool import VariablePool + + +class WorkflowNodeAndResult: + node: BaseNode + result: Optional[NodeRunResult] = None + + def __init__(self, node: BaseNode, result: Optional[NodeRunResult] = None): + self.node = node + self.result = result + + +class WorkflowRunState: + tenant_id: str + app_id: str + workflow_id: str + workflow_type: WorkflowType + user_id: str + user_from: UserFrom + invoke_from: InvokeFrom + + workflow_call_depth: int + + start_at: float + variable_pool: VariablePool + + total_tokens: int = 0 + + workflow_nodes_and_results: list[WorkflowNodeAndResult] + + class NodeRun(BaseModel): + node_id: str + iteration_node_id: str + + workflow_node_runs: list[NodeRun] + workflow_node_steps: int + + current_iteration_state: Optional[BaseIterationState] + + def __init__( + self, + workflow: Workflow, + start_at: float, + variable_pool: VariablePool, + user_id: str, + user_from: UserFrom, + invoke_from: InvokeFrom, + workflow_call_depth: int, + ): + self.workflow_id = workflow.id + self.tenant_id = workflow.tenant_id + self.app_id = workflow.app_id + self.workflow_type = WorkflowType.value_of(workflow.type) + self.user_id = user_id + self.user_from = user_from + self.invoke_from = invoke_from + self.workflow_call_depth = workflow_call_depth + + self.start_at = start_at + self.variable_pool = variable_pool + + self.total_tokens = 0 + + self.workflow_node_steps = 1 + self.workflow_node_runs = [] + self.current_iteration_state = None diff --git a/api/core/workflow/enums.py b/api/core/workflow/enums.py new file mode 100644 index 0000000000000000000000000000000000000000..9642efa1a5ca9e91e2ab232e6a6179b711bb0224 --- /dev/null +++ b/api/core/workflow/enums.py @@ -0,0 +1,16 @@ +from enum import StrEnum + + +class SystemVariableKey(StrEnum): + """ + System Variables. + """ + + QUERY = "query" + FILES = "files" + CONVERSATION_ID = "conversation_id" + USER_ID = "user_id" + DIALOGUE_COUNT = "dialogue_count" + APP_ID = "app_id" + WORKFLOW_ID = "workflow_id" + WORKFLOW_RUN_ID = "workflow_run_id" diff --git a/api/core/workflow/errors.py b/api/core/workflow/errors.py new file mode 100644 index 0000000000000000000000000000000000000000..bd4ccc1072a2a3027117d9a1317ddeea705e8773 --- /dev/null +++ b/api/core/workflow/errors.py @@ -0,0 +1,8 @@ +from core.workflow.nodes.base import BaseNode + + +class WorkflowNodeRunFailedError(Exception): + def __init__(self, node_instance: BaseNode, error: str): + self.node_instance = node_instance + self.error = error + super().__init__(f"Node {node_instance.node_data.title} run failed: {error}") diff --git a/api/core/workflow/graph_engine/__init__.py b/api/core/workflow/graph_engine/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2fee3d7fad8644fa48b6de52f13c01023f33f8be --- /dev/null +++ b/api/core/workflow/graph_engine/__init__.py @@ -0,0 +1,3 @@ +from .entities import Graph, GraphInitParams, GraphRuntimeState, RuntimeRouteState + +__all__ = ["Graph", "GraphInitParams", "GraphRuntimeState", "RuntimeRouteState"] diff --git a/api/core/workflow/graph_engine/condition_handlers/__init__.py b/api/core/workflow/graph_engine/condition_handlers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/workflow/graph_engine/condition_handlers/base_handler.py b/api/core/workflow/graph_engine/condition_handlers/base_handler.py new file mode 100644 index 0000000000000000000000000000000000000000..697392b2a3c23f26c294f9a55d4c44ce3f63cec5 --- /dev/null +++ b/api/core/workflow/graph_engine/condition_handlers/base_handler.py @@ -0,0 +1,25 @@ +from abc import ABC, abstractmethod + +from core.workflow.graph_engine.entities.graph import Graph +from core.workflow.graph_engine.entities.graph_init_params import GraphInitParams +from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState +from core.workflow.graph_engine.entities.run_condition import RunCondition +from core.workflow.graph_engine.entities.runtime_route_state import RouteNodeState + + +class RunConditionHandler(ABC): + def __init__(self, init_params: GraphInitParams, graph: Graph, condition: RunCondition): + self.init_params = init_params + self.graph = graph + self.condition = condition + + @abstractmethod + def check(self, graph_runtime_state: GraphRuntimeState, previous_route_node_state: RouteNodeState) -> bool: + """ + Check if the condition can be executed + + :param graph_runtime_state: graph runtime state + :param previous_route_node_state: previous route node state + :return: bool + """ + raise NotImplementedError diff --git a/api/core/workflow/graph_engine/condition_handlers/branch_identify_handler.py b/api/core/workflow/graph_engine/condition_handlers/branch_identify_handler.py new file mode 100644 index 0000000000000000000000000000000000000000..af695df7d84607079b59f41a95cd785645bc457a --- /dev/null +++ b/api/core/workflow/graph_engine/condition_handlers/branch_identify_handler.py @@ -0,0 +1,25 @@ +from core.workflow.graph_engine.condition_handlers.base_handler import RunConditionHandler +from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState +from core.workflow.graph_engine.entities.runtime_route_state import RouteNodeState + + +class BranchIdentifyRunConditionHandler(RunConditionHandler): + def check(self, graph_runtime_state: GraphRuntimeState, previous_route_node_state: RouteNodeState) -> bool: + """ + Check if the condition can be executed + + :param graph_runtime_state: graph runtime state + :param previous_route_node_state: previous route node state + :return: bool + """ + if not self.condition.branch_identify: + raise Exception("Branch identify is required") + + run_result = previous_route_node_state.node_run_result + if not run_result: + return False + + if not run_result.edge_source_handle: + return False + + return self.condition.branch_identify == run_result.edge_source_handle diff --git a/api/core/workflow/graph_engine/condition_handlers/condition_handler.py b/api/core/workflow/graph_engine/condition_handlers/condition_handler.py new file mode 100644 index 0000000000000000000000000000000000000000..b8470aecbd83a2b5255ee639aced3d5590939aa8 --- /dev/null +++ b/api/core/workflow/graph_engine/condition_handlers/condition_handler.py @@ -0,0 +1,27 @@ +from core.workflow.graph_engine.condition_handlers.base_handler import RunConditionHandler +from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState +from core.workflow.graph_engine.entities.runtime_route_state import RouteNodeState +from core.workflow.utils.condition.processor import ConditionProcessor + + +class ConditionRunConditionHandlerHandler(RunConditionHandler): + def check(self, graph_runtime_state: GraphRuntimeState, previous_route_node_state: RouteNodeState): + """ + Check if the condition can be executed + + :param graph_runtime_state: graph runtime state + :param previous_route_node_state: previous route node state + :return: bool + """ + if not self.condition.conditions: + return True + + # process condition + condition_processor = ConditionProcessor() + _, _, final_result = condition_processor.process_conditions( + variable_pool=graph_runtime_state.variable_pool, + conditions=self.condition.conditions, + operator="and", + ) + + return final_result diff --git a/api/core/workflow/graph_engine/condition_handlers/condition_manager.py b/api/core/workflow/graph_engine/condition_handlers/condition_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..1c9237d82fbe6887398b84393a13dc7e22b2f656 --- /dev/null +++ b/api/core/workflow/graph_engine/condition_handlers/condition_manager.py @@ -0,0 +1,25 @@ +from core.workflow.graph_engine.condition_handlers.base_handler import RunConditionHandler +from core.workflow.graph_engine.condition_handlers.branch_identify_handler import BranchIdentifyRunConditionHandler +from core.workflow.graph_engine.condition_handlers.condition_handler import ConditionRunConditionHandlerHandler +from core.workflow.graph_engine.entities.graph import Graph +from core.workflow.graph_engine.entities.graph_init_params import GraphInitParams +from core.workflow.graph_engine.entities.run_condition import RunCondition + + +class ConditionManager: + @staticmethod + def get_condition_handler( + init_params: GraphInitParams, graph: Graph, run_condition: RunCondition + ) -> RunConditionHandler: + """ + Get condition handler + + :param init_params: init params + :param graph: graph + :param run_condition: run condition + :return: condition handler + """ + if run_condition.type == "branch_identify": + return BranchIdentifyRunConditionHandler(init_params=init_params, graph=graph, condition=run_condition) + else: + return ConditionRunConditionHandlerHandler(init_params=init_params, graph=graph, condition=run_condition) diff --git a/api/core/workflow/graph_engine/entities/__init__.py b/api/core/workflow/graph_engine/entities/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6331a0b723fd507a7c7fcb451a7fa3f7cb55feb2 --- /dev/null +++ b/api/core/workflow/graph_engine/entities/__init__.py @@ -0,0 +1,6 @@ +from .graph import Graph +from .graph_init_params import GraphInitParams +from .graph_runtime_state import GraphRuntimeState +from .runtime_route_state import RuntimeRouteState + +__all__ = ["Graph", "GraphInitParams", "GraphRuntimeState", "RuntimeRouteState"] diff --git a/api/core/workflow/graph_engine/entities/event.py b/api/core/workflow/graph_engine/entities/event.py new file mode 100644 index 0000000000000000000000000000000000000000..d591b68e7e72be03d445f4ac3c6f293ce0dccf22 --- /dev/null +++ b/api/core/workflow/graph_engine/entities/event.py @@ -0,0 +1,189 @@ +from collections.abc import Mapping +from datetime import datetime +from typing import Any, Optional + +from pydantic import BaseModel, Field + +from core.workflow.graph_engine.entities.runtime_route_state import RouteNodeState +from core.workflow.nodes import NodeType +from core.workflow.nodes.base import BaseNodeData + + +class GraphEngineEvent(BaseModel): + pass + + +########################################### +# Graph Events +########################################### + + +class BaseGraphEvent(GraphEngineEvent): + pass + + +class GraphRunStartedEvent(BaseGraphEvent): + pass + + +class GraphRunSucceededEvent(BaseGraphEvent): + outputs: Optional[dict[str, Any]] = None + """outputs""" + + +class GraphRunFailedEvent(BaseGraphEvent): + error: str = Field(..., description="failed reason") + exceptions_count: int = Field(description="exception count", default=0) + + +class GraphRunPartialSucceededEvent(BaseGraphEvent): + exceptions_count: int = Field(..., description="exception count") + outputs: Optional[dict[str, Any]] = None + + +########################################### +# Node Events +########################################### + + +class BaseNodeEvent(GraphEngineEvent): + id: str = Field(..., description="node execution id") + node_id: str = Field(..., description="node id") + node_type: NodeType = Field(..., description="node type") + node_data: BaseNodeData = Field(..., description="node data") + route_node_state: RouteNodeState = Field(..., description="route node state") + parallel_id: Optional[str] = None + """parallel id if node is in parallel""" + parallel_start_node_id: Optional[str] = None + """parallel start node id if node is in parallel""" + parent_parallel_id: Optional[str] = None + """parent parallel id if node is in parallel""" + parent_parallel_start_node_id: Optional[str] = None + """parent parallel start node id if node is in parallel""" + in_iteration_id: Optional[str] = None + """iteration id if node is in iteration""" + + +class NodeRunStartedEvent(BaseNodeEvent): + predecessor_node_id: Optional[str] = None + parallel_mode_run_id: Optional[str] = None + """predecessor node id""" + + +class NodeRunStreamChunkEvent(BaseNodeEvent): + chunk_content: str = Field(..., description="chunk content") + from_variable_selector: Optional[list[str]] = None + """from variable selector""" + + +class NodeRunRetrieverResourceEvent(BaseNodeEvent): + retriever_resources: list[dict] = Field(..., description="retriever resources") + context: str = Field(..., description="context") + + +class NodeRunSucceededEvent(BaseNodeEvent): + pass + + +class NodeRunFailedEvent(BaseNodeEvent): + error: str = Field(..., description="error") + + +class NodeRunExceptionEvent(BaseNodeEvent): + error: str = Field(..., description="error") + + +class NodeInIterationFailedEvent(BaseNodeEvent): + error: str = Field(..., description="error") + + +class NodeRunRetryEvent(NodeRunStartedEvent): + error: str = Field(..., description="error") + retry_index: int = Field(..., description="which retry attempt is about to be performed") + start_at: datetime = Field(..., description="retry start time") + + +########################################### +# Parallel Branch Events +########################################### + + +class BaseParallelBranchEvent(GraphEngineEvent): + parallel_id: str = Field(..., description="parallel id") + """parallel id""" + parallel_start_node_id: str = Field(..., description="parallel start node id") + """parallel start node id""" + parent_parallel_id: Optional[str] = None + """parent parallel id if node is in parallel""" + parent_parallel_start_node_id: Optional[str] = None + """parent parallel start node id if node is in parallel""" + in_iteration_id: Optional[str] = None + """iteration id if node is in iteration""" + + +class ParallelBranchRunStartedEvent(BaseParallelBranchEvent): + pass + + +class ParallelBranchRunSucceededEvent(BaseParallelBranchEvent): + pass + + +class ParallelBranchRunFailedEvent(BaseParallelBranchEvent): + error: str = Field(..., description="failed reason") + + +########################################### +# Iteration Events +########################################### + + +class BaseIterationEvent(GraphEngineEvent): + iteration_id: str = Field(..., description="iteration node execution id") + iteration_node_id: str = Field(..., description="iteration node id") + iteration_node_type: NodeType = Field(..., description="node type, iteration or loop") + iteration_node_data: BaseNodeData = Field(..., description="node data") + parallel_id: Optional[str] = None + """parallel id if node is in parallel""" + parallel_start_node_id: Optional[str] = None + """parallel start node id if node is in parallel""" + parent_parallel_id: Optional[str] = None + """parent parallel id if node is in parallel""" + parent_parallel_start_node_id: Optional[str] = None + """parent parallel start node id if node is in parallel""" + parallel_mode_run_id: Optional[str] = None + """iteratoin run in parallel mode run id""" + + +class IterationRunStartedEvent(BaseIterationEvent): + start_at: datetime = Field(..., description="start at") + inputs: Optional[Mapping[str, Any]] = None + metadata: Optional[Mapping[str, Any]] = None + predecessor_node_id: Optional[str] = None + + +class IterationRunNextEvent(BaseIterationEvent): + index: int = Field(..., description="index") + pre_iteration_output: Optional[Any] = Field(None, description="pre iteration output") + duration: Optional[float] = Field(None, description="duration") + + +class IterationRunSucceededEvent(BaseIterationEvent): + start_at: datetime = Field(..., description="start at") + inputs: Optional[Mapping[str, Any]] = None + outputs: Optional[Mapping[str, Any]] = None + metadata: Optional[Mapping[str, Any]] = None + steps: int = 0 + iteration_duration_map: Optional[dict[str, float]] = None + + +class IterationRunFailedEvent(BaseIterationEvent): + start_at: datetime = Field(..., description="start at") + inputs: Optional[Mapping[str, Any]] = None + outputs: Optional[Mapping[str, Any]] = None + metadata: Optional[Mapping[str, Any]] = None + steps: int = 0 + error: str = Field(..., description="failed reason") + + +InNodeEvent = BaseNodeEvent | BaseParallelBranchEvent | BaseIterationEvent diff --git a/api/core/workflow/graph_engine/entities/graph.py b/api/core/workflow/graph_engine/entities/graph.py new file mode 100644 index 0000000000000000000000000000000000000000..1c6b4b6618448fb3eaf489609ac6548f6c5cecfa --- /dev/null +++ b/api/core/workflow/graph_engine/entities/graph.py @@ -0,0 +1,721 @@ +import uuid +from collections import defaultdict +from collections.abc import Mapping +from typing import Any, Optional, cast + +from pydantic import BaseModel, Field + +from configs import dify_config +from core.workflow.graph_engine.entities.run_condition import RunCondition +from core.workflow.nodes import NodeType +from core.workflow.nodes.answer.answer_stream_generate_router import AnswerStreamGeneratorRouter +from core.workflow.nodes.answer.entities import AnswerStreamGenerateRoute +from core.workflow.nodes.end.end_stream_generate_router import EndStreamGeneratorRouter +from core.workflow.nodes.end.entities import EndStreamParam + + +class GraphEdge(BaseModel): + source_node_id: str = Field(..., description="source node id") + target_node_id: str = Field(..., description="target node id") + run_condition: Optional[RunCondition] = None + """run condition""" + + +class GraphParallel(BaseModel): + id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="random uuid parallel id") + start_from_node_id: str = Field(..., description="start from node id") + parent_parallel_id: Optional[str] = None + """parent parallel id""" + parent_parallel_start_node_id: Optional[str] = None + """parent parallel start node id""" + end_to_node_id: Optional[str] = None + """end to node id""" + + +class Graph(BaseModel): + root_node_id: str = Field(..., description="root node id of the graph") + node_ids: list[str] = Field(default_factory=list, description="graph node ids") + node_id_config_mapping: dict[str, dict] = Field( + default_factory=list, description="node configs mapping (node id: node config)" + ) + edge_mapping: dict[str, list[GraphEdge]] = Field( + default_factory=dict, description="graph edge mapping (source node id: edges)" + ) + reverse_edge_mapping: dict[str, list[GraphEdge]] = Field( + default_factory=dict, description="reverse graph edge mapping (target node id: edges)" + ) + parallel_mapping: dict[str, GraphParallel] = Field( + default_factory=dict, description="graph parallel mapping (parallel id: parallel)" + ) + node_parallel_mapping: dict[str, str] = Field( + default_factory=dict, description="graph node parallel mapping (node id: parallel id)" + ) + answer_stream_generate_routes: AnswerStreamGenerateRoute = Field(..., description="answer stream generate routes") + end_stream_param: EndStreamParam = Field(..., description="end stream param") + + @classmethod + def init(cls, graph_config: Mapping[str, Any], root_node_id: Optional[str] = None) -> "Graph": + """ + Init graph + + :param graph_config: graph config + :param root_node_id: root node id + :return: graph + """ + # edge configs + edge_configs = graph_config.get("edges") + if edge_configs is None: + edge_configs = [] + # node configs + node_configs = graph_config.get("nodes") + if not node_configs: + raise ValueError("Graph must have at least one node") + + edge_configs = cast(list, edge_configs) + node_configs = cast(list, node_configs) + + # reorganize edges mapping + edge_mapping: dict[str, list[GraphEdge]] = {} + reverse_edge_mapping: dict[str, list[GraphEdge]] = {} + target_edge_ids = set() + fail_branch_source_node_id = [ + node["id"] for node in node_configs if node["data"].get("error_strategy") == "fail-branch" + ] + for edge_config in edge_configs: + source_node_id = edge_config.get("source") + if not source_node_id: + continue + + if source_node_id not in edge_mapping: + edge_mapping[source_node_id] = [] + + target_node_id = edge_config.get("target") + if not target_node_id: + continue + + if target_node_id not in reverse_edge_mapping: + reverse_edge_mapping[target_node_id] = [] + + target_edge_ids.add(target_node_id) + + # parse run condition + run_condition = None + if edge_config.get("sourceHandle"): + if ( + edge_config.get("source") in fail_branch_source_node_id + and edge_config.get("sourceHandle") != "fail-branch" + ): + run_condition = RunCondition(type="branch_identify", branch_identify="success-branch") + elif edge_config.get("sourceHandle") != "source": + run_condition = RunCondition( + type="branch_identify", branch_identify=edge_config.get("sourceHandle") + ) + + graph_edge = GraphEdge( + source_node_id=source_node_id, target_node_id=target_node_id, run_condition=run_condition + ) + + edge_mapping[source_node_id].append(graph_edge) + reverse_edge_mapping[target_node_id].append(graph_edge) + + # fetch nodes that have no predecessor node + root_node_configs = [] + all_node_id_config_mapping: dict[str, dict] = {} + for node_config in node_configs: + node_id = node_config.get("id") + if not node_id: + continue + + if node_id not in target_edge_ids: + root_node_configs.append(node_config) + + all_node_id_config_mapping[node_id] = node_config + + root_node_ids = [node_config.get("id") for node_config in root_node_configs] + + # fetch root node + if not root_node_id: + # if no root node id, use the START type node as root node + root_node_id = next( + ( + node_config.get("id") + for node_config in root_node_configs + if node_config.get("data", {}).get("type", "") == NodeType.START.value + ), + None, + ) + + if not root_node_id or root_node_id not in root_node_ids: + raise ValueError(f"Root node id {root_node_id} not found in the graph") + + # Check whether it is connected to the previous node + cls._check_connected_to_previous_node(route=[root_node_id], edge_mapping=edge_mapping) + + # fetch all node ids from root node + node_ids = [root_node_id] + cls._recursively_add_node_ids(node_ids=node_ids, edge_mapping=edge_mapping, node_id=root_node_id) + + node_id_config_mapping = {node_id: all_node_id_config_mapping[node_id] for node_id in node_ids} + + # init parallel mapping + parallel_mapping: dict[str, GraphParallel] = {} + node_parallel_mapping: dict[str, str] = {} + cls._recursively_add_parallels( + edge_mapping=edge_mapping, + reverse_edge_mapping=reverse_edge_mapping, + start_node_id=root_node_id, + parallel_mapping=parallel_mapping, + node_parallel_mapping=node_parallel_mapping, + ) + + # Check if it exceeds N layers of parallel + for parallel in parallel_mapping.values(): + if parallel.parent_parallel_id: + cls._check_exceed_parallel_limit( + parallel_mapping=parallel_mapping, + level_limit=dify_config.WORKFLOW_PARALLEL_DEPTH_LIMIT, + parent_parallel_id=parallel.parent_parallel_id, + ) + + # init answer stream generate routes + answer_stream_generate_routes = AnswerStreamGeneratorRouter.init( + node_id_config_mapping=node_id_config_mapping, reverse_edge_mapping=reverse_edge_mapping + ) + + # init end stream param + end_stream_param = EndStreamGeneratorRouter.init( + node_id_config_mapping=node_id_config_mapping, + reverse_edge_mapping=reverse_edge_mapping, + node_parallel_mapping=node_parallel_mapping, + ) + + # init graph + graph = cls( + root_node_id=root_node_id, + node_ids=node_ids, + node_id_config_mapping=node_id_config_mapping, + edge_mapping=edge_mapping, + reverse_edge_mapping=reverse_edge_mapping, + parallel_mapping=parallel_mapping, + node_parallel_mapping=node_parallel_mapping, + answer_stream_generate_routes=answer_stream_generate_routes, + end_stream_param=end_stream_param, + ) + + return graph + + def add_extra_edge( + self, source_node_id: str, target_node_id: str, run_condition: Optional[RunCondition] = None + ) -> None: + """ + Add extra edge to the graph + + :param source_node_id: source node id + :param target_node_id: target node id + :param run_condition: run condition + """ + if source_node_id not in self.node_ids or target_node_id not in self.node_ids: + return + + if source_node_id not in self.edge_mapping: + self.edge_mapping[source_node_id] = [] + + if target_node_id in [graph_edge.target_node_id for graph_edge in self.edge_mapping[source_node_id]]: + return + + graph_edge = GraphEdge( + source_node_id=source_node_id, target_node_id=target_node_id, run_condition=run_condition + ) + + self.edge_mapping[source_node_id].append(graph_edge) + + def get_leaf_node_ids(self) -> list[str]: + """ + Get leaf node ids of the graph + + :return: leaf node ids + """ + leaf_node_ids = [] + for node_id in self.node_ids: + if node_id not in self.edge_mapping or ( + len(self.edge_mapping[node_id]) == 1 + and self.edge_mapping[node_id][0].target_node_id == self.root_node_id + ): + leaf_node_ids.append(node_id) + + return leaf_node_ids + + @classmethod + def _recursively_add_node_ids( + cls, node_ids: list[str], edge_mapping: dict[str, list[GraphEdge]], node_id: str + ) -> None: + """ + Recursively add node ids + + :param node_ids: node ids + :param edge_mapping: edge mapping + :param node_id: node id + """ + for graph_edge in edge_mapping.get(node_id, []): + if graph_edge.target_node_id in node_ids: + continue + + node_ids.append(graph_edge.target_node_id) + cls._recursively_add_node_ids( + node_ids=node_ids, edge_mapping=edge_mapping, node_id=graph_edge.target_node_id + ) + + @classmethod + def _check_connected_to_previous_node(cls, route: list[str], edge_mapping: dict[str, list[GraphEdge]]) -> None: + """ + Check whether it is connected to the previous node + """ + last_node_id = route[-1] + + for graph_edge in edge_mapping.get(last_node_id, []): + if not graph_edge.target_node_id: + continue + + if graph_edge.target_node_id in route: + raise ValueError( + f"Node {graph_edge.source_node_id} is connected to the previous node, please check the graph." + ) + + new_route = route.copy() + new_route.append(graph_edge.target_node_id) + cls._check_connected_to_previous_node( + route=new_route, + edge_mapping=edge_mapping, + ) + + @classmethod + def _recursively_add_parallels( + cls, + edge_mapping: dict[str, list[GraphEdge]], + reverse_edge_mapping: dict[str, list[GraphEdge]], + start_node_id: str, + parallel_mapping: dict[str, GraphParallel], + node_parallel_mapping: dict[str, str], + parent_parallel: Optional[GraphParallel] = None, + ) -> None: + """ + Recursively add parallel ids + + :param edge_mapping: edge mapping + :param start_node_id: start from node id + :param parallel_mapping: parallel mapping + :param node_parallel_mapping: node parallel mapping + :param parent_parallel: parent parallel + """ + target_node_edges = edge_mapping.get(start_node_id, []) + parallel = None + if len(target_node_edges) > 1: + # fetch all node ids in current parallels + parallel_branch_node_ids = defaultdict(list) + condition_edge_mappings = defaultdict(list) + for graph_edge in target_node_edges: + if graph_edge.run_condition is None: + parallel_branch_node_ids["default"].append(graph_edge.target_node_id) + else: + condition_hash = graph_edge.run_condition.hash + condition_edge_mappings[condition_hash].append(graph_edge) + + for condition_hash, graph_edges in condition_edge_mappings.items(): + if len(graph_edges) > 1: + for graph_edge in graph_edges: + parallel_branch_node_ids[condition_hash].append(graph_edge.target_node_id) + + condition_parallels = {} + for condition_hash, condition_parallel_branch_node_ids in parallel_branch_node_ids.items(): + # any target node id in node_parallel_mapping + parallel = None + if condition_parallel_branch_node_ids: + parent_parallel_id = parent_parallel.id if parent_parallel else None + + parallel = GraphParallel( + start_from_node_id=start_node_id, + parent_parallel_id=parent_parallel.id if parent_parallel else None, + parent_parallel_start_node_id=parent_parallel.start_from_node_id if parent_parallel else None, + ) + parallel_mapping[parallel.id] = parallel + condition_parallels[condition_hash] = parallel + + in_branch_node_ids = cls._fetch_all_node_ids_in_parallels( + edge_mapping=edge_mapping, + reverse_edge_mapping=reverse_edge_mapping, + parallel_branch_node_ids=condition_parallel_branch_node_ids, + ) + + # collect all branches node ids + parallel_node_ids = [] + for _, node_ids in in_branch_node_ids.items(): + for node_id in node_ids: + in_parent_parallel = True + if parent_parallel_id: + in_parent_parallel = False + for parallel_node_id, parallel_id in node_parallel_mapping.items(): + if parallel_id == parent_parallel_id and parallel_node_id == node_id: + in_parent_parallel = True + break + + if in_parent_parallel: + parallel_node_ids.append(node_id) + node_parallel_mapping[node_id] = parallel.id + + outside_parallel_target_node_ids = set() + for node_id in parallel_node_ids: + if node_id == parallel.start_from_node_id: + continue + + node_edges = edge_mapping.get(node_id) + if not node_edges: + continue + + if len(node_edges) > 1: + continue + + target_node_id = node_edges[0].target_node_id + if target_node_id in parallel_node_ids: + continue + + if parent_parallel_id: + parent_parallel = parallel_mapping.get(parent_parallel_id) + if not parent_parallel: + continue + + if ( + ( + node_parallel_mapping.get(target_node_id) + and node_parallel_mapping.get(target_node_id) == parent_parallel_id + ) + or ( + parent_parallel + and parent_parallel.end_to_node_id + and target_node_id == parent_parallel.end_to_node_id + ) + or (not node_parallel_mapping.get(target_node_id) and not parent_parallel) + ): + outside_parallel_target_node_ids.add(target_node_id) + + if len(outside_parallel_target_node_ids) == 1: + if ( + parent_parallel + and parent_parallel.end_to_node_id + and parallel.end_to_node_id == parent_parallel.end_to_node_id + ): + parallel.end_to_node_id = None + else: + parallel.end_to_node_id = outside_parallel_target_node_ids.pop() + + if condition_edge_mappings: + for condition_hash, graph_edges in condition_edge_mappings.items(): + for graph_edge in graph_edges: + current_parallel = cls._get_current_parallel( + parallel_mapping=parallel_mapping, + graph_edge=graph_edge, + parallel=condition_parallels.get(condition_hash), + parent_parallel=parent_parallel, + ) + + cls._recursively_add_parallels( + edge_mapping=edge_mapping, + reverse_edge_mapping=reverse_edge_mapping, + start_node_id=graph_edge.target_node_id, + parallel_mapping=parallel_mapping, + node_parallel_mapping=node_parallel_mapping, + parent_parallel=current_parallel, + ) + else: + for graph_edge in target_node_edges: + current_parallel = cls._get_current_parallel( + parallel_mapping=parallel_mapping, + graph_edge=graph_edge, + parallel=parallel, + parent_parallel=parent_parallel, + ) + + cls._recursively_add_parallels( + edge_mapping=edge_mapping, + reverse_edge_mapping=reverse_edge_mapping, + start_node_id=graph_edge.target_node_id, + parallel_mapping=parallel_mapping, + node_parallel_mapping=node_parallel_mapping, + parent_parallel=current_parallel, + ) + else: + for graph_edge in target_node_edges: + current_parallel = cls._get_current_parallel( + parallel_mapping=parallel_mapping, + graph_edge=graph_edge, + parallel=parallel, + parent_parallel=parent_parallel, + ) + + cls._recursively_add_parallels( + edge_mapping=edge_mapping, + reverse_edge_mapping=reverse_edge_mapping, + start_node_id=graph_edge.target_node_id, + parallel_mapping=parallel_mapping, + node_parallel_mapping=node_parallel_mapping, + parent_parallel=current_parallel, + ) + + @classmethod + def _get_current_parallel( + cls, + parallel_mapping: dict[str, GraphParallel], + graph_edge: GraphEdge, + parallel: Optional[GraphParallel] = None, + parent_parallel: Optional[GraphParallel] = None, + ) -> Optional[GraphParallel]: + """ + Get current parallel + """ + current_parallel = None + if parallel: + current_parallel = parallel + elif parent_parallel: + if not parent_parallel.end_to_node_id or ( + parent_parallel.end_to_node_id and graph_edge.target_node_id != parent_parallel.end_to_node_id + ): + current_parallel = parent_parallel + else: + # fetch parent parallel's parent parallel + parent_parallel_parent_parallel_id = parent_parallel.parent_parallel_id + if parent_parallel_parent_parallel_id: + parent_parallel_parent_parallel = parallel_mapping.get(parent_parallel_parent_parallel_id) + if parent_parallel_parent_parallel and ( + not parent_parallel_parent_parallel.end_to_node_id + or ( + parent_parallel_parent_parallel.end_to_node_id + and graph_edge.target_node_id != parent_parallel_parent_parallel.end_to_node_id + ) + ): + current_parallel = parent_parallel_parent_parallel + + return current_parallel + + @classmethod + def _check_exceed_parallel_limit( + cls, + parallel_mapping: dict[str, GraphParallel], + level_limit: int, + parent_parallel_id: str, + current_level: int = 1, + ) -> None: + """ + Check if it exceeds N layers of parallel + """ + parent_parallel = parallel_mapping.get(parent_parallel_id) + if not parent_parallel: + return + + current_level += 1 + if current_level > level_limit: + raise ValueError(f"Exceeds {level_limit} layers of parallel") + + if parent_parallel.parent_parallel_id: + cls._check_exceed_parallel_limit( + parallel_mapping=parallel_mapping, + level_limit=level_limit, + parent_parallel_id=parent_parallel.parent_parallel_id, + current_level=current_level, + ) + + @classmethod + def _recursively_add_parallel_node_ids( + cls, + branch_node_ids: list[str], + edge_mapping: dict[str, list[GraphEdge]], + merge_node_id: str, + start_node_id: str, + ) -> None: + """ + Recursively add node ids + + :param branch_node_ids: in branch node ids + :param edge_mapping: edge mapping + :param merge_node_id: merge node id + :param start_node_id: start node id + """ + for graph_edge in edge_mapping.get(start_node_id, []): + if graph_edge.target_node_id != merge_node_id and graph_edge.target_node_id not in branch_node_ids: + branch_node_ids.append(graph_edge.target_node_id) + cls._recursively_add_parallel_node_ids( + branch_node_ids=branch_node_ids, + edge_mapping=edge_mapping, + merge_node_id=merge_node_id, + start_node_id=graph_edge.target_node_id, + ) + + @classmethod + def _fetch_all_node_ids_in_parallels( + cls, + edge_mapping: dict[str, list[GraphEdge]], + reverse_edge_mapping: dict[str, list[GraphEdge]], + parallel_branch_node_ids: list[str], + ) -> dict[str, list[str]]: + """ + Fetch all node ids in parallels + """ + routes_node_ids: dict[str, list[str]] = {} + for parallel_branch_node_id in parallel_branch_node_ids: + routes_node_ids[parallel_branch_node_id] = [parallel_branch_node_id] + + # fetch routes node ids + cls._recursively_fetch_routes( + edge_mapping=edge_mapping, + start_node_id=parallel_branch_node_id, + routes_node_ids=routes_node_ids[parallel_branch_node_id], + ) + + # fetch leaf node ids from routes node ids + leaf_node_ids: dict[str, list[str]] = {} + merge_branch_node_ids: dict[str, list[str]] = {} + for branch_node_id, node_ids in routes_node_ids.items(): + for node_id in node_ids: + if node_id not in edge_mapping or len(edge_mapping[node_id]) == 0: + if branch_node_id not in leaf_node_ids: + leaf_node_ids[branch_node_id] = [] + + leaf_node_ids[branch_node_id].append(node_id) + + for branch_node_id2, inner_route2 in routes_node_ids.items(): + if ( + branch_node_id != branch_node_id2 + and node_id in inner_route2 + and len(reverse_edge_mapping.get(node_id, [])) > 1 + and cls._is_node_in_routes( + reverse_edge_mapping=reverse_edge_mapping, + start_node_id=node_id, + routes_node_ids=routes_node_ids, + ) + # Exclude conditional branch nodes + and all(edge.run_condition is None for edge in reverse_edge_mapping.get(node_id, [])) + ): + if node_id not in merge_branch_node_ids: + merge_branch_node_ids[node_id] = [] + + if branch_node_id2 not in merge_branch_node_ids[node_id]: + merge_branch_node_ids[node_id].append(branch_node_id2) + + # sorted merge_branch_node_ids by branch_node_ids length desc + merge_branch_node_ids = dict(sorted(merge_branch_node_ids.items(), key=lambda x: len(x[1]), reverse=True)) + + duplicate_end_node_ids = {} + for node_id, branch_node_ids in merge_branch_node_ids.items(): + for node_id2, branch_node_ids2 in merge_branch_node_ids.items(): + if node_id != node_id2 and set(branch_node_ids) == set(branch_node_ids2): + if (node_id, node_id2) not in duplicate_end_node_ids and ( + node_id2, + node_id, + ) not in duplicate_end_node_ids: + duplicate_end_node_ids[(node_id, node_id2)] = branch_node_ids + + for (node_id, node_id2), branch_node_ids in duplicate_end_node_ids.items(): + # check which node is after + if cls._is_node2_after_node1(node1_id=node_id, node2_id=node_id2, edge_mapping=edge_mapping): + if node_id in merge_branch_node_ids and node_id2 in merge_branch_node_ids: + del merge_branch_node_ids[node_id2] + elif cls._is_node2_after_node1(node1_id=node_id2, node2_id=node_id, edge_mapping=edge_mapping): + if node_id in merge_branch_node_ids and node_id2 in merge_branch_node_ids: + del merge_branch_node_ids[node_id] + + branches_merge_node_ids: dict[str, str] = {} + for node_id, branch_node_ids in merge_branch_node_ids.items(): + if len(branch_node_ids) <= 1: + continue + + for branch_node_id in branch_node_ids: + if branch_node_id in branches_merge_node_ids: + continue + + branches_merge_node_ids[branch_node_id] = node_id + + in_branch_node_ids: dict[str, list[str]] = {} + for branch_node_id, node_ids in routes_node_ids.items(): + in_branch_node_ids[branch_node_id] = [] + if branch_node_id not in branches_merge_node_ids: + # all node ids in current branch is in this thread + in_branch_node_ids[branch_node_id].append(branch_node_id) + in_branch_node_ids[branch_node_id].extend(node_ids) + else: + merge_node_id = branches_merge_node_ids[branch_node_id] + if merge_node_id != branch_node_id: + in_branch_node_ids[branch_node_id].append(branch_node_id) + + # fetch all node ids from branch_node_id and merge_node_id + cls._recursively_add_parallel_node_ids( + branch_node_ids=in_branch_node_ids[branch_node_id], + edge_mapping=edge_mapping, + merge_node_id=merge_node_id, + start_node_id=branch_node_id, + ) + + return in_branch_node_ids + + @classmethod + def _recursively_fetch_routes( + cls, edge_mapping: dict[str, list[GraphEdge]], start_node_id: str, routes_node_ids: list[str] + ) -> None: + """ + Recursively fetch route + """ + if start_node_id not in edge_mapping: + return + + for graph_edge in edge_mapping[start_node_id]: + # find next node ids + if graph_edge.target_node_id not in routes_node_ids: + routes_node_ids.append(graph_edge.target_node_id) + + cls._recursively_fetch_routes( + edge_mapping=edge_mapping, start_node_id=graph_edge.target_node_id, routes_node_ids=routes_node_ids + ) + + @classmethod + def _is_node_in_routes( + cls, reverse_edge_mapping: dict[str, list[GraphEdge]], start_node_id: str, routes_node_ids: dict[str, list[str]] + ) -> bool: + """ + Recursively check if the node is in the routes + """ + if start_node_id not in reverse_edge_mapping: + return False + + all_routes_node_ids = set() + parallel_start_node_ids: dict[str, list[str]] = {} + for branch_node_id, node_ids in routes_node_ids.items(): + all_routes_node_ids.update(node_ids) + + if branch_node_id in reverse_edge_mapping: + for graph_edge in reverse_edge_mapping[branch_node_id]: + if graph_edge.source_node_id not in parallel_start_node_ids: + parallel_start_node_ids[graph_edge.source_node_id] = [] + + parallel_start_node_ids[graph_edge.source_node_id].append(branch_node_id) + + for _, branch_node_ids in parallel_start_node_ids.items(): + if set(branch_node_ids) == set(routes_node_ids.keys()): + return True + + return False + + @classmethod + def _is_node2_after_node1(cls, node1_id: str, node2_id: str, edge_mapping: dict[str, list[GraphEdge]]) -> bool: + """ + is node2 after node1 + """ + if node1_id not in edge_mapping: + return False + + for graph_edge in edge_mapping[node1_id]: + if graph_edge.target_node_id == node2_id: + return True + + if cls._is_node2_after_node1( + node1_id=graph_edge.target_node_id, node2_id=node2_id, edge_mapping=edge_mapping + ): + return True + + return False diff --git a/api/core/workflow/graph_engine/entities/graph_init_params.py b/api/core/workflow/graph_engine/entities/graph_init_params.py new file mode 100644 index 0000000000000000000000000000000000000000..a0ecd824f427b9dfff7618d5b22470a7dc7d74e6 --- /dev/null +++ b/api/core/workflow/graph_engine/entities/graph_init_params.py @@ -0,0 +1,21 @@ +from collections.abc import Mapping +from typing import Any + +from pydantic import BaseModel, Field + +from core.app.entities.app_invoke_entities import InvokeFrom +from models.enums import UserFrom +from models.workflow import WorkflowType + + +class GraphInitParams(BaseModel): + # init params + tenant_id: str = Field(..., description="tenant / workspace id") + app_id: str = Field(..., description="app id") + workflow_type: WorkflowType = Field(..., description="workflow type") + workflow_id: str = Field(..., description="workflow id") + graph_config: Mapping[str, Any] = Field(..., description="graph config") + user_id: str = Field(..., description="user id") + user_from: UserFrom = Field(..., description="user from, account or end-user") + invoke_from: InvokeFrom = Field(..., description="invoke from, service-api, web-app, explore or debugger") + call_depth: int = Field(..., description="call depth") diff --git a/api/core/workflow/graph_engine/entities/graph_runtime_state.py b/api/core/workflow/graph_engine/entities/graph_runtime_state.py new file mode 100644 index 0000000000000000000000000000000000000000..afc09bfac5b0c16d3933c32255e29c6b9d8f0b82 --- /dev/null +++ b/api/core/workflow/graph_engine/entities/graph_runtime_state.py @@ -0,0 +1,27 @@ +from typing import Any + +from pydantic import BaseModel, Field + +from core.model_runtime.entities.llm_entities import LLMUsage +from core.workflow.entities.variable_pool import VariablePool +from core.workflow.graph_engine.entities.runtime_route_state import RuntimeRouteState + + +class GraphRuntimeState(BaseModel): + variable_pool: VariablePool = Field(..., description="variable pool") + """variable pool""" + + start_at: float = Field(..., description="start time") + """start time""" + total_tokens: int = 0 + """total tokens""" + llm_usage: LLMUsage = LLMUsage.empty_usage() + """llm usage info""" + outputs: dict[str, Any] = {} + """outputs""" + + node_run_steps: int = 0 + """node run steps""" + + node_run_state: RuntimeRouteState = RuntimeRouteState() + """node run state""" diff --git a/api/core/workflow/graph_engine/entities/next_graph_node.py b/api/core/workflow/graph_engine/entities/next_graph_node.py new file mode 100644 index 0000000000000000000000000000000000000000..6aa4341ddfe171f9a1d998b164672e8b20b81345 --- /dev/null +++ b/api/core/workflow/graph_engine/entities/next_graph_node.py @@ -0,0 +1,13 @@ +from typing import Optional + +from pydantic import BaseModel + +from core.workflow.graph_engine.entities.graph import GraphParallel + + +class NextGraphNode(BaseModel): + node_id: str + """next node id""" + + parallel: Optional[GraphParallel] = None + """parallel""" diff --git a/api/core/workflow/graph_engine/entities/run_condition.py b/api/core/workflow/graph_engine/entities/run_condition.py new file mode 100644 index 0000000000000000000000000000000000000000..eedce8842b411efed353b9e046cd55b57213942c --- /dev/null +++ b/api/core/workflow/graph_engine/entities/run_condition.py @@ -0,0 +1,21 @@ +import hashlib +from typing import Literal, Optional + +from pydantic import BaseModel + +from core.workflow.utils.condition.entities import Condition + + +class RunCondition(BaseModel): + type: Literal["branch_identify", "condition"] + """condition type""" + + branch_identify: Optional[str] = None + """branch identify like: sourceHandle, required when type is branch_identify""" + + conditions: Optional[list[Condition]] = None + """conditions to run the node, required when type is condition""" + + @property + def hash(self) -> str: + return hashlib.sha256(self.model_dump_json().encode()).hexdigest() diff --git a/api/core/workflow/graph_engine/entities/runtime_route_state.py b/api/core/workflow/graph_engine/entities/runtime_route_state.py new file mode 100644 index 0000000000000000000000000000000000000000..7683dcc9dcd3c09afa73c144d0d6eeb877f45a59 --- /dev/null +++ b/api/core/workflow/graph_engine/entities/runtime_route_state.py @@ -0,0 +1,117 @@ +import uuid +from datetime import UTC, datetime +from enum import Enum +from typing import Optional + +from pydantic import BaseModel, Field + +from core.workflow.entities.node_entities import NodeRunResult +from models.workflow import WorkflowNodeExecutionStatus + + +class RouteNodeState(BaseModel): + class Status(Enum): + RUNNING = "running" + SUCCESS = "success" + FAILED = "failed" + PAUSED = "paused" + EXCEPTION = "exception" + + id: str = Field(default_factory=lambda: str(uuid.uuid4())) + """node state id""" + + node_id: str + """node id""" + + node_run_result: Optional[NodeRunResult] = None + """node run result""" + + status: Status = Status.RUNNING + """node status""" + + start_at: datetime + """start time""" + + paused_at: Optional[datetime] = None + """paused time""" + + finished_at: Optional[datetime] = None + """finished time""" + + failed_reason: Optional[str] = None + """failed reason""" + + paused_by: Optional[str] = None + """paused by""" + + index: int = 1 + + def set_finished(self, run_result: NodeRunResult) -> None: + """ + Node finished + + :param run_result: run result + """ + if self.status in { + RouteNodeState.Status.SUCCESS, + RouteNodeState.Status.FAILED, + RouteNodeState.Status.EXCEPTION, + }: + raise Exception(f"Route state {self.id} already finished") + + if run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED: + self.status = RouteNodeState.Status.SUCCESS + elif run_result.status == WorkflowNodeExecutionStatus.FAILED: + self.status = RouteNodeState.Status.FAILED + self.failed_reason = run_result.error + elif run_result.status == WorkflowNodeExecutionStatus.EXCEPTION: + self.status = RouteNodeState.Status.EXCEPTION + self.failed_reason = run_result.error + else: + raise Exception(f"Invalid route status {run_result.status}") + + self.node_run_result = run_result + self.finished_at = datetime.now(UTC).replace(tzinfo=None) + + +class RuntimeRouteState(BaseModel): + routes: dict[str, list[str]] = Field( + default_factory=dict, description="graph state routes (source_node_state_id: target_node_state_id)" + ) + + node_state_mapping: dict[str, RouteNodeState] = Field( + default_factory=dict, description="node state mapping (route_node_state_id: route_node_state)" + ) + + def create_node_state(self, node_id: str) -> RouteNodeState: + """ + Create node state + + :param node_id: node id + """ + state = RouteNodeState(node_id=node_id, start_at=datetime.now(UTC).replace(tzinfo=None)) + self.node_state_mapping[state.id] = state + return state + + def add_route(self, source_node_state_id: str, target_node_state_id: str) -> None: + """ + Add route to the graph state + + :param source_node_state_id: source node state id + :param target_node_state_id: target node state id + """ + if source_node_state_id not in self.routes: + self.routes[source_node_state_id] = [] + + self.routes[source_node_state_id].append(target_node_state_id) + + def get_routes_with_node_state_by_source_node_state_id(self, source_node_state_id: str) -> list[RouteNodeState]: + """ + Get routes with node state by source node id + + :param source_node_state_id: source node state id + :return: routes with node state + """ + return [ + self.node_state_mapping[target_state_id] for target_state_id in self.routes.get(source_node_state_id, []) + ] diff --git a/api/core/workflow/graph_engine/graph_engine.py b/api/core/workflow/graph_engine/graph_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..db1e01f14fda5987e3edb97aa8a19e99fe569943 --- /dev/null +++ b/api/core/workflow/graph_engine/graph_engine.py @@ -0,0 +1,911 @@ +import logging +import queue +import time +import uuid +from collections.abc import Generator, Mapping +from concurrent.futures import ThreadPoolExecutor, wait +from copy import copy, deepcopy +from datetime import UTC, datetime +from typing import Any, Optional, cast + +from flask import Flask, current_app + +from configs import dify_config +from core.app.apps.base_app_queue_manager import GenerateTaskStoppedError +from core.app.entities.app_invoke_entities import InvokeFrom +from core.workflow.entities.node_entities import NodeRunMetadataKey, NodeRunResult +from core.workflow.entities.variable_pool import VariablePool, VariableValue +from core.workflow.graph_engine.condition_handlers.condition_manager import ConditionManager +from core.workflow.graph_engine.entities.event import ( + BaseIterationEvent, + GraphEngineEvent, + GraphRunFailedEvent, + GraphRunPartialSucceededEvent, + GraphRunStartedEvent, + GraphRunSucceededEvent, + NodeRunExceptionEvent, + NodeRunFailedEvent, + NodeRunRetrieverResourceEvent, + NodeRunRetryEvent, + NodeRunStartedEvent, + NodeRunStreamChunkEvent, + NodeRunSucceededEvent, + ParallelBranchRunFailedEvent, + ParallelBranchRunStartedEvent, + ParallelBranchRunSucceededEvent, +) +from core.workflow.graph_engine.entities.graph import Graph, GraphEdge +from core.workflow.graph_engine.entities.graph_init_params import GraphInitParams +from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState +from core.workflow.graph_engine.entities.runtime_route_state import RouteNodeState +from core.workflow.nodes import NodeType +from core.workflow.nodes.answer.answer_stream_processor import AnswerStreamProcessor +from core.workflow.nodes.answer.base_stream_processor import StreamProcessor +from core.workflow.nodes.base import BaseNode +from core.workflow.nodes.base.entities import BaseNodeData +from core.workflow.nodes.end.end_stream_processor import EndStreamProcessor +from core.workflow.nodes.enums import ErrorStrategy, FailBranchSourceHandle +from core.workflow.nodes.event import RunCompletedEvent, RunRetrieverResourceEvent, RunStreamChunkEvent +from core.workflow.nodes.node_mapping import NODE_TYPE_CLASSES_MAPPING +from extensions.ext_database import db +from models.enums import UserFrom +from models.workflow import WorkflowNodeExecutionStatus, WorkflowType + +logger = logging.getLogger(__name__) + + +class GraphEngineThreadPool(ThreadPoolExecutor): + def __init__( + self, + max_workers=None, + thread_name_prefix="", + initializer=None, + initargs=(), + max_submit_count=dify_config.MAX_SUBMIT_COUNT, + ) -> None: + super().__init__(max_workers, thread_name_prefix, initializer, initargs) + self.max_submit_count = max_submit_count + self.submit_count = 0 + + def submit(self, fn, /, *args, **kwargs): + self.submit_count += 1 + self.check_is_full() + + return super().submit(fn, *args, **kwargs) + + def task_done_callback(self, future): + self.submit_count -= 1 + + def check_is_full(self) -> None: + if self.submit_count > self.max_submit_count: + raise ValueError(f"Max submit count {self.max_submit_count} of workflow thread pool reached.") + + +class GraphEngine: + workflow_thread_pool_mapping: dict[str, GraphEngineThreadPool] = {} + + def __init__( + self, + tenant_id: str, + app_id: str, + workflow_type: WorkflowType, + workflow_id: str, + user_id: str, + user_from: UserFrom, + invoke_from: InvokeFrom, + call_depth: int, + graph: Graph, + graph_config: Mapping[str, Any], + variable_pool: VariablePool, + max_execution_steps: int, + max_execution_time: int, + thread_pool_id: Optional[str] = None, + ) -> None: + thread_pool_max_submit_count = dify_config.MAX_SUBMIT_COUNT + thread_pool_max_workers = 10 + + # init thread pool + if thread_pool_id: + if thread_pool_id not in GraphEngine.workflow_thread_pool_mapping: + raise ValueError(f"Max submit count {thread_pool_max_submit_count} of workflow thread pool reached.") + + self.thread_pool_id = thread_pool_id + self.thread_pool = GraphEngine.workflow_thread_pool_mapping[thread_pool_id] + self.is_main_thread_pool = False + else: + self.thread_pool = GraphEngineThreadPool( + max_workers=thread_pool_max_workers, max_submit_count=thread_pool_max_submit_count + ) + self.thread_pool_id = str(uuid.uuid4()) + self.is_main_thread_pool = True + GraphEngine.workflow_thread_pool_mapping[self.thread_pool_id] = self.thread_pool + + self.graph = graph + self.init_params = GraphInitParams( + tenant_id=tenant_id, + app_id=app_id, + workflow_type=workflow_type, + workflow_id=workflow_id, + graph_config=graph_config, + user_id=user_id, + user_from=user_from, + invoke_from=invoke_from, + call_depth=call_depth, + ) + + self.graph_runtime_state = GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()) + + self.max_execution_steps = max_execution_steps + self.max_execution_time = max_execution_time + + def run(self) -> Generator[GraphEngineEvent, None, None]: + # trigger graph run start event + yield GraphRunStartedEvent() + handle_exceptions: list[str] = [] + stream_processor: StreamProcessor + + try: + if self.init_params.workflow_type == WorkflowType.CHAT: + stream_processor = AnswerStreamProcessor( + graph=self.graph, variable_pool=self.graph_runtime_state.variable_pool + ) + else: + stream_processor = EndStreamProcessor( + graph=self.graph, variable_pool=self.graph_runtime_state.variable_pool + ) + + # run graph + generator = stream_processor.process( + self._run(start_node_id=self.graph.root_node_id, handle_exceptions=handle_exceptions) + ) + for item in generator: + try: + yield item + if isinstance(item, NodeRunFailedEvent): + yield GraphRunFailedEvent( + error=item.route_node_state.failed_reason or "Unknown error.", + exceptions_count=len(handle_exceptions), + ) + return + elif isinstance(item, NodeRunSucceededEvent): + if item.node_type == NodeType.END: + self.graph_runtime_state.outputs = ( + dict(item.route_node_state.node_run_result.outputs) + if item.route_node_state.node_run_result + and item.route_node_state.node_run_result.outputs + else {} + ) + elif item.node_type == NodeType.ANSWER: + if "answer" not in self.graph_runtime_state.outputs: + self.graph_runtime_state.outputs["answer"] = "" + + self.graph_runtime_state.outputs["answer"] += "\n" + ( + item.route_node_state.node_run_result.outputs.get("answer", "") + if item.route_node_state.node_run_result + and item.route_node_state.node_run_result.outputs + else "" + ) + + self.graph_runtime_state.outputs["answer"] = self.graph_runtime_state.outputs[ + "answer" + ].strip() + except Exception as e: + logger.exception("Graph run failed") + yield GraphRunFailedEvent(error=str(e), exceptions_count=len(handle_exceptions)) + return + # count exceptions to determine partial success + if len(handle_exceptions) > 0: + yield GraphRunPartialSucceededEvent( + exceptions_count=len(handle_exceptions), outputs=self.graph_runtime_state.outputs + ) + else: + # trigger graph run success event + yield GraphRunSucceededEvent(outputs=self.graph_runtime_state.outputs) + self._release_thread() + except GraphRunFailedError as e: + yield GraphRunFailedEvent(error=e.error, exceptions_count=len(handle_exceptions)) + self._release_thread() + return + except Exception as e: + logger.exception("Unknown Error when graph running") + yield GraphRunFailedEvent(error=str(e), exceptions_count=len(handle_exceptions)) + self._release_thread() + raise e + + def _release_thread(self): + if self.is_main_thread_pool and self.thread_pool_id in GraphEngine.workflow_thread_pool_mapping: + del GraphEngine.workflow_thread_pool_mapping[self.thread_pool_id] + + def _run( + self, + start_node_id: str, + in_parallel_id: Optional[str] = None, + parent_parallel_id: Optional[str] = None, + parent_parallel_start_node_id: Optional[str] = None, + handle_exceptions: list[str] = [], + ) -> Generator[GraphEngineEvent, None, None]: + parallel_start_node_id = None + if in_parallel_id: + parallel_start_node_id = start_node_id + + next_node_id = start_node_id + previous_route_node_state: Optional[RouteNodeState] = None + while True: + # max steps reached + if self.graph_runtime_state.node_run_steps > self.max_execution_steps: + raise GraphRunFailedError("Max steps {} reached.".format(self.max_execution_steps)) + + # or max execution time reached + if self._is_timed_out( + start_at=self.graph_runtime_state.start_at, max_execution_time=self.max_execution_time + ): + raise GraphRunFailedError("Max execution time {}s reached.".format(self.max_execution_time)) + + # init route node state + route_node_state = self.graph_runtime_state.node_run_state.create_node_state(node_id=next_node_id) + + # get node config + node_id = route_node_state.node_id + node_config = self.graph.node_id_config_mapping.get(node_id) + if not node_config: + raise GraphRunFailedError(f"Node {node_id} config not found.") + + # convert to specific node + node_type = NodeType(node_config.get("data", {}).get("type")) + node_version = node_config.get("data", {}).get("version", "1") + node_cls = NODE_TYPE_CLASSES_MAPPING[node_type][node_version] + + previous_node_id = previous_route_node_state.node_id if previous_route_node_state else None + + # init workflow run state + node_instance = node_cls( # type: ignore + id=route_node_state.id, + config=node_config, + graph_init_params=self.init_params, + graph=self.graph, + graph_runtime_state=self.graph_runtime_state, + previous_node_id=previous_node_id, + thread_pool_id=self.thread_pool_id, + ) + node_instance = cast(BaseNode[BaseNodeData], node_instance) + try: + # run node + generator = self._run_node( + node_instance=node_instance, + route_node_state=route_node_state, + parallel_id=in_parallel_id, + parallel_start_node_id=parallel_start_node_id, + parent_parallel_id=parent_parallel_id, + parent_parallel_start_node_id=parent_parallel_start_node_id, + handle_exceptions=handle_exceptions, + ) + + for item in generator: + if isinstance(item, NodeRunStartedEvent): + self.graph_runtime_state.node_run_steps += 1 + item.route_node_state.index = self.graph_runtime_state.node_run_steps + + yield item + + self.graph_runtime_state.node_run_state.node_state_mapping[route_node_state.id] = route_node_state + + # append route + if previous_route_node_state: + self.graph_runtime_state.node_run_state.add_route( + source_node_state_id=previous_route_node_state.id, target_node_state_id=route_node_state.id + ) + except Exception as e: + route_node_state.status = RouteNodeState.Status.FAILED + route_node_state.failed_reason = str(e) + yield NodeRunFailedEvent( + error=str(e), + id=node_instance.id, + node_id=next_node_id, + node_type=node_type, + node_data=node_instance.node_data, + route_node_state=route_node_state, + parallel_id=in_parallel_id, + parallel_start_node_id=parallel_start_node_id, + parent_parallel_id=parent_parallel_id, + parent_parallel_start_node_id=parent_parallel_start_node_id, + ) + raise e + + # It may not be necessary, but it is necessary. :) + if ( + self.graph.node_id_config_mapping[next_node_id].get("data", {}).get("type", "").lower() + == NodeType.END.value + ): + break + + previous_route_node_state = route_node_state + + # get next node ids + edge_mappings = self.graph.edge_mapping.get(next_node_id) + if not edge_mappings: + break + + if len(edge_mappings) == 1: + edge = edge_mappings[0] + if ( + previous_route_node_state.status == RouteNodeState.Status.EXCEPTION + and node_instance.node_data.error_strategy == ErrorStrategy.FAIL_BRANCH + and edge.run_condition is None + ): + break + if edge.run_condition: + result = ConditionManager.get_condition_handler( + init_params=self.init_params, + graph=self.graph, + run_condition=edge.run_condition, + ).check( + graph_runtime_state=self.graph_runtime_state, + previous_route_node_state=previous_route_node_state, + ) + + if not result: + break + + next_node_id = edge.target_node_id + else: + final_node_id = None + + if any(edge.run_condition for edge in edge_mappings): + # if nodes has run conditions, get node id which branch to take based on the run condition results + condition_edge_mappings: dict[str, list[GraphEdge]] = {} + for edge in edge_mappings: + if edge.run_condition: + run_condition_hash = edge.run_condition.hash + if run_condition_hash not in condition_edge_mappings: + condition_edge_mappings[run_condition_hash] = [] + + condition_edge_mappings[run_condition_hash].append(edge) + + for _, sub_edge_mappings in condition_edge_mappings.items(): + if len(sub_edge_mappings) == 0: + continue + + edge = cast(GraphEdge, sub_edge_mappings[0]) + if edge.run_condition is None: + logger.warning(f"Edge {edge.target_node_id} run condition is None") + continue + + result = ConditionManager.get_condition_handler( + init_params=self.init_params, + graph=self.graph, + run_condition=edge.run_condition, + ).check( + graph_runtime_state=self.graph_runtime_state, + previous_route_node_state=previous_route_node_state, + ) + + if not result: + continue + + if len(sub_edge_mappings) == 1: + final_node_id = edge.target_node_id + else: + parallel_generator = self._run_parallel_branches( + edge_mappings=sub_edge_mappings, + in_parallel_id=in_parallel_id, + parallel_start_node_id=parallel_start_node_id, + handle_exceptions=handle_exceptions, + ) + + for parallel_result in parallel_generator: + if isinstance(parallel_result, str): + final_node_id = parallel_result + else: + yield parallel_result + + break + + if not final_node_id: + break + + next_node_id = final_node_id + elif ( + node_instance.node_data.error_strategy == ErrorStrategy.FAIL_BRANCH + and node_instance.should_continue_on_error + and previous_route_node_state.status == RouteNodeState.Status.EXCEPTION + ): + break + else: + parallel_generator = self._run_parallel_branches( + edge_mappings=edge_mappings, + in_parallel_id=in_parallel_id, + parallel_start_node_id=parallel_start_node_id, + handle_exceptions=handle_exceptions, + ) + + for generated_item in parallel_generator: + if isinstance(generated_item, str): + final_node_id = generated_item + else: + yield generated_item + + if not final_node_id: + break + + next_node_id = final_node_id + + if in_parallel_id and self.graph.node_parallel_mapping.get(next_node_id, "") != in_parallel_id: + break + + def _run_parallel_branches( + self, + edge_mappings: list[GraphEdge], + in_parallel_id: Optional[str] = None, + parallel_start_node_id: Optional[str] = None, + handle_exceptions: list[str] = [], + ) -> Generator[GraphEngineEvent | str, None, None]: + # if nodes has no run conditions, parallel run all nodes + parallel_id = self.graph.node_parallel_mapping.get(edge_mappings[0].target_node_id) + if not parallel_id: + node_id = edge_mappings[0].target_node_id + node_config = self.graph.node_id_config_mapping.get(node_id) + if not node_config: + raise GraphRunFailedError( + f"Node {node_id} related parallel not found or incorrectly connected to multiple parallel branches." + ) + + node_title = node_config.get("data", {}).get("title") + raise GraphRunFailedError( + f"Node {node_title} related parallel not found or incorrectly connected to multiple parallel branches." + ) + + parallel = self.graph.parallel_mapping.get(parallel_id) + if not parallel: + raise GraphRunFailedError(f"Parallel {parallel_id} not found.") + + # run parallel nodes, run in new thread and use queue to get results + q: queue.Queue = queue.Queue() + + # Create a list to store the threads + futures = [] + + # new thread + for edge in edge_mappings: + if ( + edge.target_node_id not in self.graph.node_parallel_mapping + or self.graph.node_parallel_mapping.get(edge.target_node_id, "") != parallel_id + ): + continue + + future = self.thread_pool.submit( + self._run_parallel_node, + **{ + "flask_app": current_app._get_current_object(), # type: ignore[attr-defined] + "q": q, + "parallel_id": parallel_id, + "parallel_start_node_id": edge.target_node_id, + "parent_parallel_id": in_parallel_id, + "parent_parallel_start_node_id": parallel_start_node_id, + "handle_exceptions": handle_exceptions, + }, + ) + + future.add_done_callback(self.thread_pool.task_done_callback) + + futures.append(future) + + succeeded_count = 0 + while True: + try: + event = q.get(timeout=1) + if event is None: + break + + yield event + if event.parallel_id == parallel_id: + if isinstance(event, ParallelBranchRunSucceededEvent): + succeeded_count += 1 + if succeeded_count == len(futures): + q.put(None) + + continue + elif isinstance(event, ParallelBranchRunFailedEvent): + raise GraphRunFailedError(event.error) + except queue.Empty: + continue + + # wait all threads + wait(futures) + + # get final node id + final_node_id = parallel.end_to_node_id + if final_node_id: + yield final_node_id + + def _run_parallel_node( + self, + flask_app: Flask, + q: queue.Queue, + parallel_id: str, + parallel_start_node_id: str, + parent_parallel_id: Optional[str] = None, + parent_parallel_start_node_id: Optional[str] = None, + handle_exceptions: list[str] = [], + ) -> None: + """ + Run parallel nodes + """ + with flask_app.app_context(): + try: + q.put( + ParallelBranchRunStartedEvent( + parallel_id=parallel_id, + parallel_start_node_id=parallel_start_node_id, + parent_parallel_id=parent_parallel_id, + parent_parallel_start_node_id=parent_parallel_start_node_id, + ) + ) + + # run node + generator = self._run( + start_node_id=parallel_start_node_id, + in_parallel_id=parallel_id, + parent_parallel_id=parent_parallel_id, + parent_parallel_start_node_id=parent_parallel_start_node_id, + handle_exceptions=handle_exceptions, + ) + + for item in generator: + q.put(item) + + # trigger graph run success event + q.put( + ParallelBranchRunSucceededEvent( + parallel_id=parallel_id, + parallel_start_node_id=parallel_start_node_id, + parent_parallel_id=parent_parallel_id, + parent_parallel_start_node_id=parent_parallel_start_node_id, + ) + ) + except GraphRunFailedError as e: + q.put( + ParallelBranchRunFailedEvent( + parallel_id=parallel_id, + parallel_start_node_id=parallel_start_node_id, + parent_parallel_id=parent_parallel_id, + parent_parallel_start_node_id=parent_parallel_start_node_id, + error=e.error, + ) + ) + except Exception as e: + logger.exception("Unknown Error when generating in parallel") + q.put( + ParallelBranchRunFailedEvent( + parallel_id=parallel_id, + parallel_start_node_id=parallel_start_node_id, + parent_parallel_id=parent_parallel_id, + parent_parallel_start_node_id=parent_parallel_start_node_id, + error=str(e), + ) + ) + finally: + db.session.remove() + + def _run_node( + self, + node_instance: BaseNode[BaseNodeData], + route_node_state: RouteNodeState, + parallel_id: Optional[str] = None, + parallel_start_node_id: Optional[str] = None, + parent_parallel_id: Optional[str] = None, + parent_parallel_start_node_id: Optional[str] = None, + handle_exceptions: list[str] = [], + ) -> Generator[GraphEngineEvent, None, None]: + """ + Run node + """ + # trigger node run start event + yield NodeRunStartedEvent( + id=node_instance.id, + node_id=node_instance.node_id, + node_type=node_instance.node_type, + node_data=node_instance.node_data, + route_node_state=route_node_state, + predecessor_node_id=node_instance.previous_node_id, + parallel_id=parallel_id, + parallel_start_node_id=parallel_start_node_id, + parent_parallel_id=parent_parallel_id, + parent_parallel_start_node_id=parent_parallel_start_node_id, + ) + + db.session.close() + max_retries = node_instance.node_data.retry_config.max_retries + retry_interval = node_instance.node_data.retry_config.retry_interval_seconds + retries = 0 + should_continue_retry = True + while should_continue_retry and retries <= max_retries: + try: + # run node + retry_start_at = datetime.now(UTC).replace(tzinfo=None) + generator = node_instance.run() + for item in generator: + if isinstance(item, GraphEngineEvent): + if isinstance(item, BaseIterationEvent): + # add parallel info to iteration event + item.parallel_id = parallel_id + item.parallel_start_node_id = parallel_start_node_id + item.parent_parallel_id = parent_parallel_id + item.parent_parallel_start_node_id = parent_parallel_start_node_id + + yield item + else: + if isinstance(item, RunCompletedEvent): + run_result = item.run_result + if run_result.status == WorkflowNodeExecutionStatus.FAILED: + if ( + retries == max_retries + and node_instance.node_type == NodeType.HTTP_REQUEST + and run_result.outputs + and not node_instance.should_continue_on_error + ): + run_result.status = WorkflowNodeExecutionStatus.SUCCEEDED + if node_instance.should_retry and retries < max_retries: + retries += 1 + route_node_state.node_run_result = run_result + yield NodeRunRetryEvent( + id=node_instance.id, + node_id=node_instance.node_id, + node_type=node_instance.node_type, + node_data=node_instance.node_data, + route_node_state=route_node_state, + predecessor_node_id=node_instance.previous_node_id, + parallel_id=parallel_id, + parallel_start_node_id=parallel_start_node_id, + parent_parallel_id=parent_parallel_id, + parent_parallel_start_node_id=parent_parallel_start_node_id, + error=run_result.error or "Unknown error", + retry_index=retries, + start_at=retry_start_at, + ) + time.sleep(retry_interval) + continue + route_node_state.set_finished(run_result=run_result) + + if run_result.status == WorkflowNodeExecutionStatus.FAILED: + if node_instance.should_continue_on_error: + # if run failed, handle error + run_result = self._handle_continue_on_error( + node_instance, + item.run_result, + self.graph_runtime_state.variable_pool, + handle_exceptions=handle_exceptions, + ) + route_node_state.node_run_result = run_result + route_node_state.status = RouteNodeState.Status.EXCEPTION + if run_result.outputs: + for variable_key, variable_value in run_result.outputs.items(): + # append variables to variable pool recursively + self._append_variables_recursively( + node_id=node_instance.node_id, + variable_key_list=[variable_key], + variable_value=variable_value, + ) + yield NodeRunExceptionEvent( + error=run_result.error or "System Error", + id=node_instance.id, + node_id=node_instance.node_id, + node_type=node_instance.node_type, + node_data=node_instance.node_data, + route_node_state=route_node_state, + parallel_id=parallel_id, + parallel_start_node_id=parallel_start_node_id, + parent_parallel_id=parent_parallel_id, + parent_parallel_start_node_id=parent_parallel_start_node_id, + ) + should_continue_retry = False + else: + yield NodeRunFailedEvent( + error=route_node_state.failed_reason or "Unknown error.", + id=node_instance.id, + node_id=node_instance.node_id, + node_type=node_instance.node_type, + node_data=node_instance.node_data, + route_node_state=route_node_state, + parallel_id=parallel_id, + parallel_start_node_id=parallel_start_node_id, + parent_parallel_id=parent_parallel_id, + parent_parallel_start_node_id=parent_parallel_start_node_id, + ) + should_continue_retry = False + elif run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED: + if node_instance.should_continue_on_error and self.graph.edge_mapping.get( + node_instance.node_id + ): + run_result.edge_source_handle = FailBranchSourceHandle.SUCCESS + if run_result.metadata and run_result.metadata.get(NodeRunMetadataKey.TOTAL_TOKENS): + # plus state total_tokens + self.graph_runtime_state.total_tokens += int( + run_result.metadata.get(NodeRunMetadataKey.TOTAL_TOKENS) # type: ignore[arg-type] + ) + + if run_result.llm_usage: + # use the latest usage + self.graph_runtime_state.llm_usage += run_result.llm_usage + + # append node output variables to variable pool + if run_result.outputs: + for variable_key, variable_value in run_result.outputs.items(): + # append variables to variable pool recursively + self._append_variables_recursively( + node_id=node_instance.node_id, + variable_key_list=[variable_key], + variable_value=variable_value, + ) + + # When setting metadata, convert to dict first + if not run_result.metadata: + run_result.metadata = {} + + if parallel_id and parallel_start_node_id: + metadata_dict = dict(run_result.metadata) + metadata_dict[NodeRunMetadataKey.PARALLEL_ID] = parallel_id + metadata_dict[NodeRunMetadataKey.PARALLEL_START_NODE_ID] = parallel_start_node_id + if parent_parallel_id and parent_parallel_start_node_id: + metadata_dict[NodeRunMetadataKey.PARENT_PARALLEL_ID] = parent_parallel_id + metadata_dict[NodeRunMetadataKey.PARENT_PARALLEL_START_NODE_ID] = ( + parent_parallel_start_node_id + ) + run_result.metadata = metadata_dict + + yield NodeRunSucceededEvent( + id=node_instance.id, + node_id=node_instance.node_id, + node_type=node_instance.node_type, + node_data=node_instance.node_data, + route_node_state=route_node_state, + parallel_id=parallel_id, + parallel_start_node_id=parallel_start_node_id, + parent_parallel_id=parent_parallel_id, + parent_parallel_start_node_id=parent_parallel_start_node_id, + ) + should_continue_retry = False + + break + elif isinstance(item, RunStreamChunkEvent): + yield NodeRunStreamChunkEvent( + id=node_instance.id, + node_id=node_instance.node_id, + node_type=node_instance.node_type, + node_data=node_instance.node_data, + chunk_content=item.chunk_content, + from_variable_selector=item.from_variable_selector, + route_node_state=route_node_state, + parallel_id=parallel_id, + parallel_start_node_id=parallel_start_node_id, + parent_parallel_id=parent_parallel_id, + parent_parallel_start_node_id=parent_parallel_start_node_id, + ) + elif isinstance(item, RunRetrieverResourceEvent): + yield NodeRunRetrieverResourceEvent( + id=node_instance.id, + node_id=node_instance.node_id, + node_type=node_instance.node_type, + node_data=node_instance.node_data, + retriever_resources=item.retriever_resources, + context=item.context, + route_node_state=route_node_state, + parallel_id=parallel_id, + parallel_start_node_id=parallel_start_node_id, + parent_parallel_id=parent_parallel_id, + parent_parallel_start_node_id=parent_parallel_start_node_id, + ) + except GenerateTaskStoppedError: + # trigger node run failed event + route_node_state.status = RouteNodeState.Status.FAILED + route_node_state.failed_reason = "Workflow stopped." + yield NodeRunFailedEvent( + error="Workflow stopped.", + id=node_instance.id, + node_id=node_instance.node_id, + node_type=node_instance.node_type, + node_data=node_instance.node_data, + route_node_state=route_node_state, + parallel_id=parallel_id, + parallel_start_node_id=parallel_start_node_id, + parent_parallel_id=parent_parallel_id, + parent_parallel_start_node_id=parent_parallel_start_node_id, + ) + return + except Exception as e: + logger.exception(f"Node {node_instance.node_data.title} run failed") + raise e + finally: + db.session.close() + + def _append_variables_recursively(self, node_id: str, variable_key_list: list[str], variable_value: VariableValue): + """ + Append variables recursively + :param node_id: node id + :param variable_key_list: variable key list + :param variable_value: variable value + :return: + """ + self.graph_runtime_state.variable_pool.add([node_id] + variable_key_list, variable_value) + + # if variable_value is a dict, then recursively append variables + if isinstance(variable_value, dict): + for key, value in variable_value.items(): + # construct new key list + new_key_list = variable_key_list + [key] + self._append_variables_recursively( + node_id=node_id, variable_key_list=new_key_list, variable_value=value + ) + + def _is_timed_out(self, start_at: float, max_execution_time: int) -> bool: + """ + Check timeout + :param start_at: start time + :param max_execution_time: max execution time + :return: + """ + return time.perf_counter() - start_at > max_execution_time + + def create_copy(self): + """ + create a graph engine copy + :return: with a new variable pool instance of graph engine + """ + new_instance = copy(self) + new_instance.graph_runtime_state = copy(self.graph_runtime_state) + new_instance.graph_runtime_state.variable_pool = deepcopy(self.graph_runtime_state.variable_pool) + return new_instance + + def _handle_continue_on_error( + self, + node_instance: BaseNode[BaseNodeData], + error_result: NodeRunResult, + variable_pool: VariablePool, + handle_exceptions: list[str] = [], + ) -> NodeRunResult: + """ + handle continue on error when self._should_continue_on_error is True + + + :param error_result (NodeRunResult): error run result + :param variable_pool (VariablePool): variable pool + :return: excption run result + """ + # add error message and error type to variable pool + variable_pool.add([node_instance.node_id, "error_message"], error_result.error) + variable_pool.add([node_instance.node_id, "error_type"], error_result.error_type) + # add error message to handle_exceptions + handle_exceptions.append(error_result.error or "") + node_error_args: dict[str, Any] = { + "status": WorkflowNodeExecutionStatus.EXCEPTION, + "error": error_result.error, + "inputs": error_result.inputs, + "metadata": { + NodeRunMetadataKey.ERROR_STRATEGY: node_instance.node_data.error_strategy, + }, + } + + if node_instance.node_data.error_strategy is ErrorStrategy.DEFAULT_VALUE: + return NodeRunResult( + **node_error_args, + outputs={ + **node_instance.node_data.default_value_dict, + "error_message": error_result.error, + "error_type": error_result.error_type, + }, + ) + elif node_instance.node_data.error_strategy is ErrorStrategy.FAIL_BRANCH: + if self.graph.edge_mapping.get(node_instance.node_id): + node_error_args["edge_source_handle"] = FailBranchSourceHandle.FAILED + return NodeRunResult( + **node_error_args, + outputs={ + "error_message": error_result.error, + "error_type": error_result.error_type, + }, + ) + return error_result + + +class GraphRunFailedError(Exception): + def __init__(self, error: str): + self.error = error diff --git a/api/core/workflow/nodes/__init__.py b/api/core/workflow/nodes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6101fcf9afd982dbebd1c64bb2ca530d1caa1ddd --- /dev/null +++ b/api/core/workflow/nodes/__init__.py @@ -0,0 +1,3 @@ +from .enums import NodeType + +__all__ = ["NodeType"] diff --git a/api/core/workflow/nodes/answer/__init__.py b/api/core/workflow/nodes/answer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ee7676c7e48418097b094c9db5e419fd8ca678a4 --- /dev/null +++ b/api/core/workflow/nodes/answer/__init__.py @@ -0,0 +1,4 @@ +from .answer_node import AnswerNode +from .entities import AnswerStreamGenerateRoute + +__all__ = ["AnswerNode", "AnswerStreamGenerateRoute"] diff --git a/api/core/workflow/nodes/answer/answer_node.py b/api/core/workflow/nodes/answer/answer_node.py new file mode 100644 index 0000000000000000000000000000000000000000..520cbdbb6051154749f602c9a06a84b7e470a22c --- /dev/null +++ b/api/core/workflow/nodes/answer/answer_node.py @@ -0,0 +1,72 @@ +from collections.abc import Mapping, Sequence +from typing import Any, cast + +from core.variables import ArrayFileSegment, FileSegment +from core.workflow.entities.node_entities import NodeRunResult +from core.workflow.nodes.answer.answer_stream_generate_router import AnswerStreamGeneratorRouter +from core.workflow.nodes.answer.entities import ( + AnswerNodeData, + GenerateRouteChunk, + TextGenerateRouteChunk, + VarGenerateRouteChunk, +) +from core.workflow.nodes.base import BaseNode +from core.workflow.nodes.enums import NodeType +from core.workflow.utils.variable_template_parser import VariableTemplateParser +from models.workflow import WorkflowNodeExecutionStatus + + +class AnswerNode(BaseNode[AnswerNodeData]): + _node_data_cls = AnswerNodeData + _node_type: NodeType = NodeType.ANSWER + + def _run(self) -> NodeRunResult: + """ + Run node + :return: + """ + # generate routes + generate_routes = AnswerStreamGeneratorRouter.extract_generate_route_from_node_data(self.node_data) + + answer = "" + files = [] + for part in generate_routes: + if part.type == GenerateRouteChunk.ChunkType.VAR: + part = cast(VarGenerateRouteChunk, part) + value_selector = part.value_selector + variable = self.graph_runtime_state.variable_pool.get(value_selector) + if variable: + if isinstance(variable, FileSegment): + files.append(variable.value) + elif isinstance(variable, ArrayFileSegment): + files.extend(variable.value) + answer += variable.markdown + else: + part = cast(TextGenerateRouteChunk, part) + answer += part.text + + return NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED, outputs={"answer": answer, "files": files}) + + @classmethod + def _extract_variable_selector_to_variable_mapping( + cls, + *, + graph_config: Mapping[str, Any], + node_id: str, + node_data: AnswerNodeData, + ) -> Mapping[str, Sequence[str]]: + """ + Extract variable selector to variable mapping + :param graph_config: graph config + :param node_id: node id + :param node_data: node data + :return: + """ + variable_template_parser = VariableTemplateParser(template=node_data.answer) + variable_selectors = variable_template_parser.extract_variable_selectors() + + variable_mapping = {} + for variable_selector in variable_selectors: + variable_mapping[node_id + "." + variable_selector.variable] = variable_selector.value_selector + + return variable_mapping diff --git a/api/core/workflow/nodes/answer/answer_stream_generate_router.py b/api/core/workflow/nodes/answer/answer_stream_generate_router.py new file mode 100644 index 0000000000000000000000000000000000000000..7d652d39f70ef4031e99369a2f539453e9c02f43 --- /dev/null +++ b/api/core/workflow/nodes/answer/answer_stream_generate_router.py @@ -0,0 +1,173 @@ +from core.prompt.utils.prompt_template_parser import PromptTemplateParser +from core.workflow.nodes.answer.entities import ( + AnswerNodeData, + AnswerStreamGenerateRoute, + GenerateRouteChunk, + TextGenerateRouteChunk, + VarGenerateRouteChunk, +) +from core.workflow.nodes.enums import ErrorStrategy, NodeType +from core.workflow.utils.variable_template_parser import VariableTemplateParser + + +class AnswerStreamGeneratorRouter: + @classmethod + def init( + cls, + node_id_config_mapping: dict[str, dict], + reverse_edge_mapping: dict[str, list["GraphEdge"]], # type: ignore[name-defined] + ) -> AnswerStreamGenerateRoute: + """ + Get stream generate routes. + :return: + """ + # parse stream output node value selectors of answer nodes + answer_generate_route: dict[str, list[GenerateRouteChunk]] = {} + for answer_node_id, node_config in node_id_config_mapping.items(): + if node_config.get("data", {}).get("type") != NodeType.ANSWER.value: + continue + + # get generate route for stream output + generate_route = cls._extract_generate_route_selectors(node_config) + answer_generate_route[answer_node_id] = generate_route + + # fetch answer dependencies + answer_node_ids = list(answer_generate_route.keys()) + answer_dependencies = cls._fetch_answers_dependencies( + answer_node_ids=answer_node_ids, + reverse_edge_mapping=reverse_edge_mapping, + node_id_config_mapping=node_id_config_mapping, + ) + + return AnswerStreamGenerateRoute( + answer_generate_route=answer_generate_route, answer_dependencies=answer_dependencies + ) + + @classmethod + def extract_generate_route_from_node_data(cls, node_data: AnswerNodeData) -> list[GenerateRouteChunk]: + """ + Extract generate route from node data + :param node_data: node data object + :return: + """ + variable_template_parser = VariableTemplateParser(template=node_data.answer) + variable_selectors = variable_template_parser.extract_variable_selectors() + + value_selector_mapping = { + variable_selector.variable: variable_selector.value_selector for variable_selector in variable_selectors + } + + variable_keys = list(value_selector_mapping.keys()) + + # format answer template + template_parser = PromptTemplateParser(template=node_data.answer, with_variable_tmpl=True) + template_variable_keys = template_parser.variable_keys + + # Take the intersection of variable_keys and template_variable_keys + variable_keys = list(set(variable_keys) & set(template_variable_keys)) + + template = node_data.answer + for var in variable_keys: + template = template.replace(f"{{{{{var}}}}}", f"Ω{{{{{var}}}}}Ω") + + generate_routes: list[GenerateRouteChunk] = [] + for part in template.split("Ω"): + if part: + if cls._is_variable(part, variable_keys): + var_key = part.replace("Ω", "").replace("{{", "").replace("}}", "") + value_selector = value_selector_mapping[var_key] + generate_routes.append(VarGenerateRouteChunk(value_selector=value_selector)) + else: + generate_routes.append(TextGenerateRouteChunk(text=part)) + + return generate_routes + + @classmethod + def _extract_generate_route_selectors(cls, config: dict) -> list[GenerateRouteChunk]: + """ + Extract generate route selectors + :param config: node config + :return: + """ + node_data = AnswerNodeData(**config.get("data", {})) + return cls.extract_generate_route_from_node_data(node_data) + + @classmethod + def _is_variable(cls, part, variable_keys): + cleaned_part = part.replace("{{", "").replace("}}", "") + return part.startswith("{{") and cleaned_part in variable_keys + + @classmethod + def _fetch_answers_dependencies( + cls, + answer_node_ids: list[str], + reverse_edge_mapping: dict[str, list["GraphEdge"]], # type: ignore[name-defined] + node_id_config_mapping: dict[str, dict], + ) -> dict[str, list[str]]: + """ + Fetch answer dependencies + :param answer_node_ids: answer node ids + :param reverse_edge_mapping: reverse edge mapping + :param node_id_config_mapping: node id config mapping + :return: + """ + answer_dependencies: dict[str, list[str]] = {} + for answer_node_id in answer_node_ids: + if answer_dependencies.get(answer_node_id) is None: + answer_dependencies[answer_node_id] = [] + + cls._recursive_fetch_answer_dependencies( + current_node_id=answer_node_id, + answer_node_id=answer_node_id, + node_id_config_mapping=node_id_config_mapping, + reverse_edge_mapping=reverse_edge_mapping, + answer_dependencies=answer_dependencies, + ) + + return answer_dependencies + + @classmethod + def _recursive_fetch_answer_dependencies( + cls, + current_node_id: str, + answer_node_id: str, + node_id_config_mapping: dict[str, dict], + reverse_edge_mapping: dict[str, list["GraphEdge"]], # type: ignore[name-defined] + answer_dependencies: dict[str, list[str]], + ) -> None: + """ + Recursive fetch answer dependencies + :param current_node_id: current node id + :param answer_node_id: answer node id + :param node_id_config_mapping: node id config mapping + :param reverse_edge_mapping: reverse edge mapping + :param answer_dependencies: answer dependencies + :return: + """ + reverse_edges = reverse_edge_mapping.get(current_node_id, []) + for edge in reverse_edges: + source_node_id = edge.source_node_id + if source_node_id not in node_id_config_mapping: + continue + source_node_type = node_id_config_mapping[source_node_id].get("data", {}).get("type") + source_node_data = node_id_config_mapping[source_node_id].get("data", {}) + if ( + source_node_type + in { + NodeType.ANSWER, + NodeType.IF_ELSE, + NodeType.QUESTION_CLASSIFIER, + NodeType.ITERATION, + NodeType.VARIABLE_ASSIGNER, + } + or source_node_data.get("error_strategy") == ErrorStrategy.FAIL_BRANCH + ): + answer_dependencies[answer_node_id].append(source_node_id) + else: + cls._recursive_fetch_answer_dependencies( + current_node_id=source_node_id, + answer_node_id=answer_node_id, + node_id_config_mapping=node_id_config_mapping, + reverse_edge_mapping=reverse_edge_mapping, + answer_dependencies=answer_dependencies, + ) diff --git a/api/core/workflow/nodes/answer/answer_stream_processor.py b/api/core/workflow/nodes/answer/answer_stream_processor.py new file mode 100644 index 0000000000000000000000000000000000000000..40213bd151f7afca78eb0d7eace46b6fbe86358c --- /dev/null +++ b/api/core/workflow/nodes/answer/answer_stream_processor.py @@ -0,0 +1,221 @@ +import logging +from collections.abc import Generator +from typing import cast + +from core.file import FILE_MODEL_IDENTITY, File +from core.workflow.entities.variable_pool import VariablePool +from core.workflow.graph_engine.entities.event import ( + GraphEngineEvent, + NodeRunExceptionEvent, + NodeRunStartedEvent, + NodeRunStreamChunkEvent, + NodeRunSucceededEvent, +) +from core.workflow.graph_engine.entities.graph import Graph +from core.workflow.nodes.answer.base_stream_processor import StreamProcessor +from core.workflow.nodes.answer.entities import GenerateRouteChunk, TextGenerateRouteChunk, VarGenerateRouteChunk + +logger = logging.getLogger(__name__) + + +class AnswerStreamProcessor(StreamProcessor): + def __init__(self, graph: Graph, variable_pool: VariablePool) -> None: + super().__init__(graph, variable_pool) + self.generate_routes = graph.answer_stream_generate_routes + self.route_position = {} + for answer_node_id in self.generate_routes.answer_generate_route: + self.route_position[answer_node_id] = 0 + self.current_stream_chunk_generating_node_ids: dict[str, list[str]] = {} + + def process(self, generator: Generator[GraphEngineEvent, None, None]) -> Generator[GraphEngineEvent, None, None]: + for event in generator: + if isinstance(event, NodeRunStartedEvent): + if event.route_node_state.node_id == self.graph.root_node_id and not self.rest_node_ids: + self.reset() + + yield event + elif isinstance(event, NodeRunStreamChunkEvent): + if event.in_iteration_id: + yield event + continue + + if event.route_node_state.node_id in self.current_stream_chunk_generating_node_ids: + stream_out_answer_node_ids = self.current_stream_chunk_generating_node_ids[ + event.route_node_state.node_id + ] + else: + stream_out_answer_node_ids = self._get_stream_out_answer_node_ids(event) + self.current_stream_chunk_generating_node_ids[event.route_node_state.node_id] = ( + stream_out_answer_node_ids + ) + + for _ in stream_out_answer_node_ids: + yield event + elif isinstance(event, NodeRunSucceededEvent | NodeRunExceptionEvent): + yield event + if event.route_node_state.node_id in self.current_stream_chunk_generating_node_ids: + # update self.route_position after all stream event finished + for answer_node_id in self.current_stream_chunk_generating_node_ids[event.route_node_state.node_id]: + self.route_position[answer_node_id] += 1 + + del self.current_stream_chunk_generating_node_ids[event.route_node_state.node_id] + + self._remove_unreachable_nodes(event) + + # generate stream outputs + yield from self._generate_stream_outputs_when_node_finished(cast(NodeRunSucceededEvent, event)) + else: + yield event + + def reset(self) -> None: + self.route_position = {} + for answer_node_id, route_chunks in self.generate_routes.answer_generate_route.items(): + self.route_position[answer_node_id] = 0 + self.rest_node_ids = self.graph.node_ids.copy() + self.current_stream_chunk_generating_node_ids = {} + + def _generate_stream_outputs_when_node_finished( + self, event: NodeRunSucceededEvent + ) -> Generator[GraphEngineEvent, None, None]: + """ + Generate stream outputs. + :param event: node run succeeded event + :return: + """ + for answer_node_id, position in self.route_position.items(): + # all depends on answer node id not in rest node ids + if event.route_node_state.node_id != answer_node_id and ( + answer_node_id not in self.rest_node_ids + or not all( + dep_id not in self.rest_node_ids + for dep_id in self.generate_routes.answer_dependencies[answer_node_id] + ) + ): + continue + + route_position = self.route_position[answer_node_id] + route_chunks = self.generate_routes.answer_generate_route[answer_node_id][route_position:] + + for route_chunk in route_chunks: + if route_chunk.type == GenerateRouteChunk.ChunkType.TEXT: + route_chunk = cast(TextGenerateRouteChunk, route_chunk) + yield NodeRunStreamChunkEvent( + id=event.id, + node_id=event.node_id, + node_type=event.node_type, + node_data=event.node_data, + chunk_content=route_chunk.text, + route_node_state=event.route_node_state, + parallel_id=event.parallel_id, + parallel_start_node_id=event.parallel_start_node_id, + from_variable_selector=[answer_node_id, "answer"], + ) + else: + route_chunk = cast(VarGenerateRouteChunk, route_chunk) + value_selector = route_chunk.value_selector + if not value_selector: + break + + value = self.variable_pool.get(value_selector) + + if value is None: + break + + text = value.markdown + + if text: + yield NodeRunStreamChunkEvent( + id=event.id, + node_id=event.node_id, + node_type=event.node_type, + node_data=event.node_data, + chunk_content=text, + from_variable_selector=list(value_selector), + route_node_state=event.route_node_state, + parallel_id=event.parallel_id, + parallel_start_node_id=event.parallel_start_node_id, + ) + + self.route_position[answer_node_id] += 1 + + def _get_stream_out_answer_node_ids(self, event: NodeRunStreamChunkEvent) -> list[str]: + """ + Is stream out support + :param event: queue text chunk event + :return: + """ + if not event.from_variable_selector: + return [] + + stream_output_value_selector = event.from_variable_selector + if not stream_output_value_selector: + return [] + + stream_out_answer_node_ids = [] + for answer_node_id, route_position in self.route_position.items(): + if answer_node_id not in self.rest_node_ids: + continue + + # all depends on answer node id not in rest node ids + if all( + dep_id not in self.rest_node_ids for dep_id in self.generate_routes.answer_dependencies[answer_node_id] + ): + if route_position >= len(self.generate_routes.answer_generate_route[answer_node_id]): + continue + + route_chunk = self.generate_routes.answer_generate_route[answer_node_id][route_position] + + if route_chunk.type != GenerateRouteChunk.ChunkType.VAR: + continue + + route_chunk = cast(VarGenerateRouteChunk, route_chunk) + value_selector = route_chunk.value_selector + + # check chunk node id is before current node id or equal to current node id + if value_selector != stream_output_value_selector: + continue + + stream_out_answer_node_ids.append(answer_node_id) + + return stream_out_answer_node_ids + + @classmethod + def _fetch_files_from_variable_value(cls, value: dict | list) -> list[dict]: + """ + Fetch files from variable value + :param value: variable value + :return: + """ + if not value: + return [] + + files = [] + if isinstance(value, list): + for item in value: + file_var = cls._get_file_var_from_value(item) + if file_var: + files.append(file_var) + elif isinstance(value, dict): + file_var = cls._get_file_var_from_value(value) + if file_var: + files.append(file_var) + + return files + + @classmethod + def _get_file_var_from_value(cls, value: dict | list): + """ + Get file var from value + :param value: variable value + :return: + """ + if not value: + return None + + if isinstance(value, dict): + if "dify_model_identity" in value and value["dify_model_identity"] == FILE_MODEL_IDENTITY: + return value + elif isinstance(value, File): + return value.to_dict() + + return None diff --git a/api/core/workflow/nodes/answer/base_stream_processor.py b/api/core/workflow/nodes/answer/base_stream_processor.py new file mode 100644 index 0000000000000000000000000000000000000000..4759356ae124a435935267454a8530288c76171f --- /dev/null +++ b/api/core/workflow/nodes/answer/base_stream_processor.py @@ -0,0 +1,95 @@ +import logging +from abc import ABC, abstractmethod +from collections.abc import Generator +from typing import Optional + +from core.workflow.entities.variable_pool import VariablePool +from core.workflow.graph_engine.entities.event import GraphEngineEvent, NodeRunExceptionEvent, NodeRunSucceededEvent +from core.workflow.graph_engine.entities.graph import Graph + +logger = logging.getLogger(__name__) + + +class StreamProcessor(ABC): + def __init__(self, graph: Graph, variable_pool: VariablePool) -> None: + self.graph = graph + self.variable_pool = variable_pool + self.rest_node_ids = graph.node_ids.copy() + + @abstractmethod + def process(self, generator: Generator[GraphEngineEvent, None, None]) -> Generator[GraphEngineEvent, None, None]: + raise NotImplementedError + + def _remove_unreachable_nodes(self, event: NodeRunSucceededEvent | NodeRunExceptionEvent) -> None: + finished_node_id = event.route_node_state.node_id + if finished_node_id not in self.rest_node_ids: + return + + # remove finished node id + self.rest_node_ids.remove(finished_node_id) + + run_result = event.route_node_state.node_run_result + if not run_result: + return + + if run_result.edge_source_handle: + reachable_node_ids: list[str] = [] + unreachable_first_node_ids: list[str] = [] + if finished_node_id not in self.graph.edge_mapping: + logger.warning(f"node {finished_node_id} has no edge mapping") + return + for edge in self.graph.edge_mapping[finished_node_id]: + if ( + edge.run_condition + and edge.run_condition.branch_identify + and run_result.edge_source_handle == edge.run_condition.branch_identify + ): + # remove unreachable nodes + # FIXME: because of the code branch can combine directly, so for answer node + # we remove the node maybe shortcut the answer node, so comment this code for now + # there is not effect on the answer node and the workflow, when we have a better solution + # we can open this code. Issues: #11542 #9560 #10638 #10564 + # ids = self._fetch_node_ids_in_reachable_branch(edge.target_node_id) + # if "answer" in ids: + # continue + # else: + # reachable_node_ids.extend(ids) + + # The branch_identify parameter is added to ensure that + # only nodes in the correct logical branch are included. + ids = self._fetch_node_ids_in_reachable_branch(edge.target_node_id, run_result.edge_source_handle) + reachable_node_ids.extend(ids) + else: + unreachable_first_node_ids.append(edge.target_node_id) + + for node_id in unreachable_first_node_ids: + self._remove_node_ids_in_unreachable_branch(node_id, reachable_node_ids) + + def _fetch_node_ids_in_reachable_branch(self, node_id: str, branch_identify: Optional[str] = None) -> list[str]: + node_ids = [] + for edge in self.graph.edge_mapping.get(node_id, []): + if edge.target_node_id == self.graph.root_node_id: + continue + + # Only follow edges that match the branch_identify or have no run_condition + if edge.run_condition and edge.run_condition.branch_identify: + if not branch_identify or edge.run_condition.branch_identify != branch_identify: + continue + + node_ids.append(edge.target_node_id) + node_ids.extend(self._fetch_node_ids_in_reachable_branch(edge.target_node_id, branch_identify)) + return node_ids + + def _remove_node_ids_in_unreachable_branch(self, node_id: str, reachable_node_ids: list[str]) -> None: + """ + remove target node ids until merge + """ + if node_id not in self.rest_node_ids: + return + + self.rest_node_ids.remove(node_id) + for edge in self.graph.edge_mapping.get(node_id, []): + if edge.target_node_id in reachable_node_ids: + continue + + self._remove_node_ids_in_unreachable_branch(edge.target_node_id, reachable_node_ids) diff --git a/api/core/workflow/nodes/answer/entities.py b/api/core/workflow/nodes/answer/entities.py new file mode 100644 index 0000000000000000000000000000000000000000..a05cc44c99428e70746f2f625fecaf5954c09b98 --- /dev/null +++ b/api/core/workflow/nodes/answer/entities.py @@ -0,0 +1,65 @@ +from collections.abc import Sequence +from enum import Enum + +from pydantic import BaseModel, Field + +from core.workflow.nodes.base import BaseNodeData + + +class AnswerNodeData(BaseNodeData): + """ + Answer Node Data. + """ + + answer: str = Field(..., description="answer template string") + + +class GenerateRouteChunk(BaseModel): + """ + Generate Route Chunk. + """ + + class ChunkType(Enum): + VAR = "var" + TEXT = "text" + + type: ChunkType = Field(..., description="generate route chunk type") + + +class VarGenerateRouteChunk(GenerateRouteChunk): + """ + Var Generate Route Chunk. + """ + + type: GenerateRouteChunk.ChunkType = GenerateRouteChunk.ChunkType.VAR + """generate route chunk type""" + value_selector: Sequence[str] = Field(..., description="value selector") + + +class TextGenerateRouteChunk(GenerateRouteChunk): + """ + Text Generate Route Chunk. + """ + + type: GenerateRouteChunk.ChunkType = GenerateRouteChunk.ChunkType.TEXT + """generate route chunk type""" + text: str = Field(..., description="text") + + +class AnswerNodeDoubleLink(BaseModel): + node_id: str = Field(..., description="node id") + source_node_ids: list[str] = Field(..., description="source node ids") + target_node_ids: list[str] = Field(..., description="target node ids") + + +class AnswerStreamGenerateRoute(BaseModel): + """ + AnswerStreamGenerateRoute entity + """ + + answer_dependencies: dict[str, list[str]] = Field( + ..., description="answer dependencies (answer node id -> dependent answer node ids)" + ) + answer_generate_route: dict[str, list[GenerateRouteChunk]] = Field( + ..., description="answer generate route (answer node id -> generate route chunks)" + ) diff --git a/api/core/workflow/nodes/base/__init__.py b/api/core/workflow/nodes/base/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..72d6392d4e01ea677d50c2133d1040e32fc4c8fc --- /dev/null +++ b/api/core/workflow/nodes/base/__init__.py @@ -0,0 +1,4 @@ +from .entities import BaseIterationNodeData, BaseIterationState, BaseNodeData +from .node import BaseNode + +__all__ = ["BaseIterationNodeData", "BaseIterationState", "BaseNode", "BaseNodeData"] diff --git a/api/core/workflow/nodes/base/entities.py b/api/core/workflow/nodes/base/entities.py new file mode 100644 index 0000000000000000000000000000000000000000..6bf8899f5d698b84558e8344bdb68d3bb28ee816 --- /dev/null +++ b/api/core/workflow/nodes/base/entities.py @@ -0,0 +1,149 @@ +import json +from abc import ABC +from enum import StrEnum +from typing import Any, Optional, Union + +from pydantic import BaseModel, model_validator + +from core.workflow.nodes.base.exc import DefaultValueTypeError +from core.workflow.nodes.enums import ErrorStrategy + + +class DefaultValueType(StrEnum): + STRING = "string" + NUMBER = "number" + OBJECT = "object" + ARRAY_NUMBER = "array[number]" + ARRAY_STRING = "array[string]" + ARRAY_OBJECT = "array[object]" + ARRAY_FILES = "array[file]" + + +NumberType = Union[int, float] + + +class DefaultValue(BaseModel): + value: Any + type: DefaultValueType + key: str + + @staticmethod + def _parse_json(value: str) -> Any: + """Unified JSON parsing handler""" + try: + return json.loads(value) + except json.JSONDecodeError: + raise DefaultValueTypeError(f"Invalid JSON format for value: {value}") + + @staticmethod + def _validate_array(value: Any, element_type: DefaultValueType) -> bool: + """Unified array type validation""" + # FIXME, type ignore here for do not find the reason mypy complain, if find the root cause, please fix it + return isinstance(value, list) and all(isinstance(x, element_type) for x in value) # type: ignore + + @staticmethod + def _convert_number(value: str) -> float: + """Unified number conversion handler""" + try: + return float(value) + except ValueError: + raise DefaultValueTypeError(f"Cannot convert to number: {value}") + + @model_validator(mode="after") + def validate_value_type(self) -> "DefaultValue": + if self.type is None: + raise DefaultValueTypeError("type field is required") + + # Type validation configuration + type_validators = { + DefaultValueType.STRING: { + "type": str, + "converter": lambda x: x, + }, + DefaultValueType.NUMBER: { + "type": NumberType, + "converter": self._convert_number, + }, + DefaultValueType.OBJECT: { + "type": dict, + "converter": self._parse_json, + }, + DefaultValueType.ARRAY_NUMBER: { + "type": list, + "element_type": NumberType, + "converter": self._parse_json, + }, + DefaultValueType.ARRAY_STRING: { + "type": list, + "element_type": str, + "converter": self._parse_json, + }, + DefaultValueType.ARRAY_OBJECT: { + "type": list, + "element_type": dict, + "converter": self._parse_json, + }, + } + + validator: dict[str, Any] = type_validators.get(self.type, {}) + if not validator: + if self.type == DefaultValueType.ARRAY_FILES: + # Handle files type + return self + raise DefaultValueTypeError(f"Unsupported type: {self.type}") + + # Handle string input cases + if isinstance(self.value, str) and self.type != DefaultValueType.STRING: + self.value = validator["converter"](self.value) + + # Validate base type + if not isinstance(self.value, validator["type"]): + raise DefaultValueTypeError(f"Value must be {validator['type'].__name__} type for {self.value}") + + # Validate array element types + if validator["type"] == list and not self._validate_array(self.value, validator["element_type"]): + raise DefaultValueTypeError(f"All elements must be {validator['element_type'].__name__} for {self.value}") + + return self + + +class RetryConfig(BaseModel): + """node retry config""" + + max_retries: int = 0 # max retry times + retry_interval: int = 0 # retry interval in milliseconds + retry_enabled: bool = False # whether retry is enabled + + @property + def retry_interval_seconds(self) -> float: + return self.retry_interval / 1000 + + +class BaseNodeData(ABC, BaseModel): + title: str + desc: Optional[str] = None + error_strategy: Optional[ErrorStrategy] = None + default_value: Optional[list[DefaultValue]] = None + version: str = "1" + retry_config: RetryConfig = RetryConfig() + + @property + def default_value_dict(self): + if self.default_value: + return {item.key: item.value for item in self.default_value} + return {} + + +class BaseIterationNodeData(BaseNodeData): + start_node_id: Optional[str] = None + + +class BaseIterationState(BaseModel): + iteration_node_id: str + index: int + inputs: dict + + class MetaData(BaseModel): + pass + + metadata: MetaData diff --git a/api/core/workflow/nodes/base/exc.py b/api/core/workflow/nodes/base/exc.py new file mode 100644 index 0000000000000000000000000000000000000000..aeecf406403e6d7abe510eb6a558637199a8f0f5 --- /dev/null +++ b/api/core/workflow/nodes/base/exc.py @@ -0,0 +1,10 @@ +class BaseNodeError(ValueError): + """Base class for node errors.""" + + pass + + +class DefaultValueTypeError(BaseNodeError): + """Raised when the default value type is invalid.""" + + pass diff --git a/api/core/workflow/nodes/base/node.py b/api/core/workflow/nodes/base/node.py new file mode 100644 index 0000000000000000000000000000000000000000..b799e7426616e7bae14ea21adbbcd5dbbeec5ab0 --- /dev/null +++ b/api/core/workflow/nodes/base/node.py @@ -0,0 +1,158 @@ +import logging +from abc import abstractmethod +from collections.abc import Generator, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar, Union, cast + +from core.workflow.entities.node_entities import NodeRunResult +from core.workflow.nodes.enums import CONTINUE_ON_ERROR_NODE_TYPE, RETRY_ON_ERROR_NODE_TYPE, NodeType +from core.workflow.nodes.event import NodeEvent, RunCompletedEvent +from models.workflow import WorkflowNodeExecutionStatus + +from .entities import BaseNodeData + +if TYPE_CHECKING: + from core.workflow.graph_engine.entities.event import InNodeEvent + from core.workflow.graph_engine.entities.graph import Graph + from core.workflow.graph_engine.entities.graph_init_params import GraphInitParams + from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState + +logger = logging.getLogger(__name__) + +GenericNodeData = TypeVar("GenericNodeData", bound=BaseNodeData) + + +class BaseNode(Generic[GenericNodeData]): + _node_data_cls: type[BaseNodeData] + _node_type: NodeType + + def __init__( + self, + id: str, + config: Mapping[str, Any], + graph_init_params: "GraphInitParams", + graph: "Graph", + graph_runtime_state: "GraphRuntimeState", + previous_node_id: Optional[str] = None, + thread_pool_id: Optional[str] = None, + ) -> None: + self.id = id + self.tenant_id = graph_init_params.tenant_id + self.app_id = graph_init_params.app_id + self.workflow_type = graph_init_params.workflow_type + self.workflow_id = graph_init_params.workflow_id + self.graph_config = graph_init_params.graph_config + self.user_id = graph_init_params.user_id + self.user_from = graph_init_params.user_from + self.invoke_from = graph_init_params.invoke_from + self.workflow_call_depth = graph_init_params.call_depth + self.graph = graph + self.graph_runtime_state = graph_runtime_state + self.previous_node_id = previous_node_id + self.thread_pool_id = thread_pool_id + + node_id = config.get("id") + if not node_id: + raise ValueError("Node ID is required.") + + self.node_id = node_id + + node_data = self._node_data_cls.model_validate(config.get("data", {})) + self.node_data = cast(GenericNodeData, node_data) + + @abstractmethod + def _run(self) -> NodeRunResult | Generator[Union[NodeEvent, "InNodeEvent"], None, None]: + """ + Run node + :return: + """ + raise NotImplementedError + + def run(self) -> Generator[Union[NodeEvent, "InNodeEvent"], None, None]: + try: + result = self._run() + except Exception as e: + logger.exception(f"Node {self.node_id} failed to run") + result = NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + error=str(e), + error_type="WorkflowNodeError", + ) + + if isinstance(result, NodeRunResult): + yield RunCompletedEvent(run_result=result) + else: + yield from result + + @classmethod + def extract_variable_selector_to_variable_mapping( + cls, + *, + graph_config: Mapping[str, Any], + config: Mapping[str, Any], + ) -> Mapping[str, Sequence[str]]: + """ + Extract variable selector to variable mapping + :param graph_config: graph config + :param config: node config + :return: + """ + node_id = config.get("id") + if not node_id: + raise ValueError("Node ID is required when extracting variable selector to variable mapping.") + + node_data = cls._node_data_cls(**config.get("data", {})) + return cls._extract_variable_selector_to_variable_mapping( + graph_config=graph_config, node_id=node_id, node_data=cast(GenericNodeData, node_data) + ) + + @classmethod + def _extract_variable_selector_to_variable_mapping( + cls, + *, + graph_config: Mapping[str, Any], + node_id: str, + node_data: GenericNodeData, + ) -> Mapping[str, Sequence[str]]: + """ + Extract variable selector to variable mapping + :param graph_config: graph config + :param node_id: node id + :param node_data: node data + :return: + """ + return {} + + @classmethod + def get_default_config(cls, filters: Optional[dict] = None) -> dict: + """ + Get default config of node. + :param filters: filter by node config parameters. + :return: + """ + return {} + + @property + def node_type(self) -> NodeType: + """ + Get node type + :return: + """ + return self._node_type + + @property + def should_continue_on_error(self) -> bool: + """judge if should continue on error + + Returns: + bool: if should continue on error + """ + return self.node_data.error_strategy is not None and self.node_type in CONTINUE_ON_ERROR_NODE_TYPE + + @property + def should_retry(self) -> bool: + """judge if should retry + + Returns: + bool: if should retry + """ + return self.node_data.retry_config.retry_enabled and self.node_type in RETRY_ON_ERROR_NODE_TYPE diff --git a/api/core/workflow/nodes/code/__init__.py b/api/core/workflow/nodes/code/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8c6dcc7fccbf2ed77c857bd011f2186332d74195 --- /dev/null +++ b/api/core/workflow/nodes/code/__init__.py @@ -0,0 +1,3 @@ +from .code_node import CodeNode + +__all__ = ["CodeNode"] diff --git a/api/core/workflow/nodes/code/code_node.py b/api/core/workflow/nodes/code/code_node.py new file mode 100644 index 0000000000000000000000000000000000000000..2f82bf8c382b55b6fec4f7a4e57b977432698ba6 --- /dev/null +++ b/api/core/workflow/nodes/code/code_node.py @@ -0,0 +1,332 @@ +from collections.abc import Mapping, Sequence +from typing import Any, Optional + +from configs import dify_config +from core.helper.code_executor.code_executor import CodeExecutionError, CodeExecutor, CodeLanguage +from core.helper.code_executor.code_node_provider import CodeNodeProvider +from core.helper.code_executor.javascript.javascript_code_provider import JavascriptCodeProvider +from core.helper.code_executor.python3.python3_code_provider import Python3CodeProvider +from core.workflow.entities.node_entities import NodeRunResult +from core.workflow.nodes.base import BaseNode +from core.workflow.nodes.code.entities import CodeNodeData +from core.workflow.nodes.enums import NodeType +from models.workflow import WorkflowNodeExecutionStatus + +from .exc import ( + CodeNodeError, + DepthLimitError, + OutputValidationError, +) + + +class CodeNode(BaseNode[CodeNodeData]): + _node_data_cls = CodeNodeData + _node_type = NodeType.CODE + + @classmethod + def get_default_config(cls, filters: Optional[dict] = None) -> dict: + """ + Get default config of node. + :param filters: filter by node config parameters. + :return: + """ + code_language = CodeLanguage.PYTHON3 + if filters: + code_language = filters.get("code_language", CodeLanguage.PYTHON3) + + providers: list[type[CodeNodeProvider]] = [Python3CodeProvider, JavascriptCodeProvider] + code_provider: type[CodeNodeProvider] = next(p for p in providers if p.is_accept_language(code_language)) + + return code_provider.get_default_config() + + def _run(self) -> NodeRunResult: + # Get code language + code_language = self.node_data.code_language + code = self.node_data.code + + # Get variables + variables = {} + for variable_selector in self.node_data.variables: + variable_name = variable_selector.variable + variable = self.graph_runtime_state.variable_pool.get(variable_selector.value_selector) + variables[variable_name] = variable.to_object() if variable else None + # Run code + try: + result = CodeExecutor.execute_workflow_code_template( + language=code_language, + code=code, + inputs=variables, + ) + + # Transform result + result = self._transform_result(result=result, output_schema=self.node_data.outputs) + except (CodeExecutionError, CodeNodeError) as e: + return NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, inputs=variables, error=str(e), error_type=type(e).__name__ + ) + + return NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED, inputs=variables, outputs=result) + + def _check_string(self, value: str | None, variable: str) -> str | None: + """ + Check string + :param value: value + :param variable: variable + :return: + """ + if value is None: + return None + if not isinstance(value, str): + raise OutputValidationError(f"Output variable `{variable}` must be a string") + + if len(value) > dify_config.CODE_MAX_STRING_LENGTH: + raise OutputValidationError( + f"The length of output variable `{variable}` must be" + f" less than {dify_config.CODE_MAX_STRING_LENGTH} characters" + ) + + return value.replace("\x00", "") + + def _check_number(self, value: int | float | None, variable: str) -> int | float | None: + """ + Check number + :param value: value + :param variable: variable + :return: + """ + if value is None: + return None + if not isinstance(value, int | float): + raise OutputValidationError(f"Output variable `{variable}` must be a number") + + if value > dify_config.CODE_MAX_NUMBER or value < dify_config.CODE_MIN_NUMBER: + raise OutputValidationError( + f"Output variable `{variable}` is out of range," + f" it must be between {dify_config.CODE_MIN_NUMBER} and {dify_config.CODE_MAX_NUMBER}." + ) + + if isinstance(value, float): + # raise error if precision is too high + if len(str(value).split(".")[1]) > dify_config.CODE_MAX_PRECISION: + raise OutputValidationError( + f"Output variable `{variable}` has too high precision," + f" it must be less than {dify_config.CODE_MAX_PRECISION} digits." + ) + + return value + + def _transform_result( + self, + result: Mapping[str, Any], + output_schema: Optional[dict[str, CodeNodeData.Output]], + prefix: str = "", + depth: int = 1, + ): + if depth > dify_config.CODE_MAX_DEPTH: + raise DepthLimitError(f"Depth limit ${dify_config.CODE_MAX_DEPTH} reached, object too deep.") + + transformed_result: dict[str, Any] = {} + if output_schema is None: + # validate output thought instance type + for output_name, output_value in result.items(): + if isinstance(output_value, dict): + self._transform_result( + result=output_value, + output_schema=None, + prefix=f"{prefix}.{output_name}" if prefix else output_name, + depth=depth + 1, + ) + elif isinstance(output_value, int | float): + self._check_number( + value=output_value, variable=f"{prefix}.{output_name}" if prefix else output_name + ) + elif isinstance(output_value, str): + self._check_string( + value=output_value, variable=f"{prefix}.{output_name}" if prefix else output_name + ) + elif isinstance(output_value, list): + first_element = output_value[0] if len(output_value) > 0 else None + if first_element is not None: + if isinstance(first_element, int | float) and all( + value is None or isinstance(value, int | float) for value in output_value + ): + for i, value in enumerate(output_value): + self._check_number( + value=value, + variable=f"{prefix}.{output_name}[{i}]" if prefix else f"{output_name}[{i}]", + ) + elif isinstance(first_element, str) and all( + value is None or isinstance(value, str) for value in output_value + ): + for i, value in enumerate(output_value): + self._check_string( + value=value, + variable=f"{prefix}.{output_name}[{i}]" if prefix else f"{output_name}[{i}]", + ) + elif isinstance(first_element, dict) and all( + value is None or isinstance(value, dict) for value in output_value + ): + for i, value in enumerate(output_value): + if value is not None: + self._transform_result( + result=value, + output_schema=None, + prefix=f"{prefix}.{output_name}[{i}]" if prefix else f"{output_name}[{i}]", + depth=depth + 1, + ) + else: + raise OutputValidationError( + f"Output {prefix}.{output_name} is not a valid array." + f" make sure all elements are of the same type." + ) + elif output_value is None: + pass + else: + raise OutputValidationError(f"Output {prefix}.{output_name} is not a valid type.") + + return result + + parameters_validated = {} + for output_name, output_config in output_schema.items(): + dot = "." if prefix else "" + if output_name not in result: + raise OutputValidationError(f"Output {prefix}{dot}{output_name} is missing.") + + if output_config.type == "object": + # check if output is object + if not isinstance(result.get(output_name), dict): + if isinstance(result.get(output_name), type(None)): + transformed_result[output_name] = None + else: + raise OutputValidationError( + f"Output {prefix}{dot}{output_name} is not an object," + f" got {type(result.get(output_name))} instead." + ) + else: + transformed_result[output_name] = self._transform_result( + result=result[output_name], + output_schema=output_config.children, + prefix=f"{prefix}.{output_name}", + depth=depth + 1, + ) + elif output_config.type == "number": + # check if number available + transformed_result[output_name] = self._check_number( + value=result[output_name], variable=f"{prefix}{dot}{output_name}" + ) + elif output_config.type == "string": + # check if string available + transformed_result[output_name] = self._check_string( + value=result[output_name], + variable=f"{prefix}{dot}{output_name}", + ) + elif output_config.type == "array[number]": + # check if array of number available + if not isinstance(result[output_name], list): + if isinstance(result[output_name], type(None)): + transformed_result[output_name] = None + else: + raise OutputValidationError( + f"Output {prefix}{dot}{output_name} is not an array," + f" got {type(result.get(output_name))} instead." + ) + else: + if len(result[output_name]) > dify_config.CODE_MAX_NUMBER_ARRAY_LENGTH: + raise OutputValidationError( + f"The length of output variable `{prefix}{dot}{output_name}` must be" + f" less than {dify_config.CODE_MAX_NUMBER_ARRAY_LENGTH} elements." + ) + + transformed_result[output_name] = [ + self._check_number(value=value, variable=f"{prefix}{dot}{output_name}[{i}]") + for i, value in enumerate(result[output_name]) + ] + elif output_config.type == "array[string]": + # check if array of string available + if not isinstance(result[output_name], list): + if isinstance(result[output_name], type(None)): + transformed_result[output_name] = None + else: + raise OutputValidationError( + f"Output {prefix}{dot}{output_name} is not an array," + f" got {type(result.get(output_name))} instead." + ) + else: + if len(result[output_name]) > dify_config.CODE_MAX_STRING_ARRAY_LENGTH: + raise OutputValidationError( + f"The length of output variable `{prefix}{dot}{output_name}` must be" + f" less than {dify_config.CODE_MAX_STRING_ARRAY_LENGTH} elements." + ) + + transformed_result[output_name] = [ + self._check_string(value=value, variable=f"{prefix}{dot}{output_name}[{i}]") + for i, value in enumerate(result[output_name]) + ] + elif output_config.type == "array[object]": + # check if array of object available + if not isinstance(result[output_name], list): + if isinstance(result[output_name], type(None)): + transformed_result[output_name] = None + else: + raise OutputValidationError( + f"Output {prefix}{dot}{output_name} is not an array," + f" got {type(result.get(output_name))} instead." + ) + else: + if len(result[output_name]) > dify_config.CODE_MAX_OBJECT_ARRAY_LENGTH: + raise OutputValidationError( + f"The length of output variable `{prefix}{dot}{output_name}` must be" + f" less than {dify_config.CODE_MAX_OBJECT_ARRAY_LENGTH} elements." + ) + + for i, value in enumerate(result[output_name]): + if not isinstance(value, dict): + if value is None: + pass + else: + raise OutputValidationError( + f"Output {prefix}{dot}{output_name}[{i}] is not an object," + f" got {type(value)} instead at index {i}." + ) + + transformed_result[output_name] = [ + None + if value is None + else self._transform_result( + result=value, + output_schema=output_config.children, + prefix=f"{prefix}{dot}{output_name}[{i}]", + depth=depth + 1, + ) + for i, value in enumerate(result[output_name]) + ] + else: + raise OutputValidationError(f"Output type {output_config.type} is not supported.") + + parameters_validated[output_name] = True + + # check if all output parameters are validated + if len(parameters_validated) != len(result): + raise CodeNodeError("Not all output parameters are validated.") + + return transformed_result + + @classmethod + def _extract_variable_selector_to_variable_mapping( + cls, + *, + graph_config: Mapping[str, Any], + node_id: str, + node_data: CodeNodeData, + ) -> Mapping[str, Sequence[str]]: + """ + Extract variable selector to variable mapping + :param graph_config: graph config + :param node_id: node id + :param node_data: node data + :return: + """ + return { + node_id + "." + variable_selector.variable: variable_selector.value_selector + for variable_selector in node_data.variables + } diff --git a/api/core/workflow/nodes/code/entities.py b/api/core/workflow/nodes/code/entities.py new file mode 100644 index 0000000000000000000000000000000000000000..a45403588832109a685a16fa0dee333b434c45f9 --- /dev/null +++ b/api/core/workflow/nodes/code/entities.py @@ -0,0 +1,27 @@ +from typing import Literal, Optional + +from pydantic import BaseModel + +from core.helper.code_executor.code_executor import CodeLanguage +from core.workflow.entities.variable_entities import VariableSelector +from core.workflow.nodes.base import BaseNodeData + + +class CodeNodeData(BaseNodeData): + """ + Code Node Data. + """ + + class Output(BaseModel): + type: Literal["string", "number", "object", "array[string]", "array[number]", "array[object]"] + children: Optional[dict[str, "CodeNodeData.Output"]] = None + + class Dependency(BaseModel): + name: str + version: str + + variables: list[VariableSelector] + code_language: Literal[CodeLanguage.PYTHON3, CodeLanguage.JAVASCRIPT] + code: str + outputs: dict[str, Output] + dependencies: Optional[list[Dependency]] = None diff --git a/api/core/workflow/nodes/code/exc.py b/api/core/workflow/nodes/code/exc.py new file mode 100644 index 0000000000000000000000000000000000000000..d6334fd554cde57a62f3429d242ee3caf5bc1e73 --- /dev/null +++ b/api/core/workflow/nodes/code/exc.py @@ -0,0 +1,16 @@ +class CodeNodeError(ValueError): + """Base class for code node errors.""" + + pass + + +class OutputValidationError(CodeNodeError): + """Raised when there is an output validation error.""" + + pass + + +class DepthLimitError(CodeNodeError): + """Raised when the depth limit is reached.""" + + pass diff --git a/api/core/workflow/nodes/document_extractor/__init__.py b/api/core/workflow/nodes/document_extractor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3cc5fae18745f973dbe4a8e496e95ab1dce434e8 --- /dev/null +++ b/api/core/workflow/nodes/document_extractor/__init__.py @@ -0,0 +1,4 @@ +from .entities import DocumentExtractorNodeData +from .node import DocumentExtractorNode + +__all__ = ["DocumentExtractorNode", "DocumentExtractorNodeData"] diff --git a/api/core/workflow/nodes/document_extractor/entities.py b/api/core/workflow/nodes/document_extractor/entities.py new file mode 100644 index 0000000000000000000000000000000000000000..7e9ffaa889b988c521d1211eddec865b8b093aa3 --- /dev/null +++ b/api/core/workflow/nodes/document_extractor/entities.py @@ -0,0 +1,7 @@ +from collections.abc import Sequence + +from core.workflow.nodes.base import BaseNodeData + + +class DocumentExtractorNodeData(BaseNodeData): + variable_selector: Sequence[str] diff --git a/api/core/workflow/nodes/document_extractor/exc.py b/api/core/workflow/nodes/document_extractor/exc.py new file mode 100644 index 0000000000000000000000000000000000000000..5caf00ebc5f1c6054295e85f9033f63c89a080b0 --- /dev/null +++ b/api/core/workflow/nodes/document_extractor/exc.py @@ -0,0 +1,14 @@ +class DocumentExtractorError(ValueError): + """Base exception for errors related to the DocumentExtractorNode.""" + + +class FileDownloadError(DocumentExtractorError): + """Exception raised when there's an error downloading a file.""" + + +class UnsupportedFileTypeError(DocumentExtractorError): + """Exception raised when trying to extract text from an unsupported file type.""" + + +class TextExtractionError(DocumentExtractorError): + """Exception raised when there's an error during text extraction from a file.""" diff --git a/api/core/workflow/nodes/document_extractor/node.py b/api/core/workflow/nodes/document_extractor/node.py new file mode 100644 index 0000000000000000000000000000000000000000..c0d8c6409982e6ad5825d25f1c7ccaef73fed317 --- /dev/null +++ b/api/core/workflow/nodes/document_extractor/node.py @@ -0,0 +1,397 @@ +import csv +import io +import json +import logging +import operator +import os +import tempfile +from collections.abc import Mapping, Sequence +from typing import Any, cast + +import docx +import pandas as pd +import pypdfium2 # type: ignore +import yaml # type: ignore +from docx.table import Table +from docx.text.paragraph import Paragraph + +from configs import dify_config +from core.file import File, FileTransferMethod, file_manager +from core.helper import ssrf_proxy +from core.variables import ArrayFileSegment +from core.variables.segments import FileSegment +from core.workflow.entities.node_entities import NodeRunResult +from core.workflow.nodes.base import BaseNode +from core.workflow.nodes.enums import NodeType +from models.workflow import WorkflowNodeExecutionStatus + +from .entities import DocumentExtractorNodeData +from .exc import DocumentExtractorError, FileDownloadError, TextExtractionError, UnsupportedFileTypeError + +logger = logging.getLogger(__name__) + + +class DocumentExtractorNode(BaseNode[DocumentExtractorNodeData]): + """ + Extracts text content from various file types. + Supports plain text, PDF, and DOC/DOCX files. + """ + + _node_data_cls = DocumentExtractorNodeData + _node_type = NodeType.DOCUMENT_EXTRACTOR + + def _run(self): + variable_selector = self.node_data.variable_selector + variable = self.graph_runtime_state.variable_pool.get(variable_selector) + + if variable is None: + error_message = f"File variable not found for selector: {variable_selector}" + return NodeRunResult(status=WorkflowNodeExecutionStatus.FAILED, error=error_message) + if variable.value and not isinstance(variable, ArrayFileSegment | FileSegment): + error_message = f"Variable {variable_selector} is not an ArrayFileSegment" + return NodeRunResult(status=WorkflowNodeExecutionStatus.FAILED, error=error_message) + + value = variable.value + inputs = {"variable_selector": variable_selector} + process_data = {"documents": value if isinstance(value, list) else [value]} + + try: + if isinstance(value, list): + extracted_text_list = list(map(_extract_text_from_file, value)) + return NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + inputs=inputs, + process_data=process_data, + outputs={"text": extracted_text_list}, + ) + elif isinstance(value, File): + extracted_text = _extract_text_from_file(value) + return NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + inputs=inputs, + process_data=process_data, + outputs={"text": extracted_text}, + ) + else: + raise DocumentExtractorError(f"Unsupported variable type: {type(value)}") + except DocumentExtractorError as e: + return NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + error=str(e), + inputs=inputs, + process_data=process_data, + ) + + @classmethod + def _extract_variable_selector_to_variable_mapping( + cls, + *, + graph_config: Mapping[str, Any], + node_id: str, + node_data: DocumentExtractorNodeData, + ) -> Mapping[str, Sequence[str]]: + """ + Extract variable selector to variable mapping + :param graph_config: graph config + :param node_id: node id + :param node_data: node data + :return: + """ + return {node_id + ".files": node_data.variable_selector} + + +def _extract_text_by_mime_type(*, file_content: bytes, mime_type: str) -> str: + """Extract text from a file based on its MIME type.""" + match mime_type: + case "text/plain" | "text/html" | "text/htm" | "text/markdown" | "text/xml": + return _extract_text_from_plain_text(file_content) + case "application/pdf": + return _extract_text_from_pdf(file_content) + case "application/vnd.openxmlformats-officedocument.wordprocessingml.document" | "application/msword": + return _extract_text_from_doc(file_content) + case "text/csv": + return _extract_text_from_csv(file_content) + case "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" | "application/vnd.ms-excel": + return _extract_text_from_excel(file_content) + case "application/vnd.ms-powerpoint": + return _extract_text_from_ppt(file_content) + case "application/vnd.openxmlformats-officedocument.presentationml.presentation": + return _extract_text_from_pptx(file_content) + case "application/epub+zip": + return _extract_text_from_epub(file_content) + case "message/rfc822": + return _extract_text_from_eml(file_content) + case "application/vnd.ms-outlook": + return _extract_text_from_msg(file_content) + case "application/json": + return _extract_text_from_json(file_content) + case "application/x-yaml" | "text/yaml": + return _extract_text_from_yaml(file_content) + case _: + raise UnsupportedFileTypeError(f"Unsupported MIME type: {mime_type}") + + +def _extract_text_by_file_extension(*, file_content: bytes, file_extension: str) -> str: + """Extract text from a file based on its file extension.""" + match file_extension: + case ".txt" | ".markdown" | ".md" | ".html" | ".htm" | ".xml" | ".vtt": + return _extract_text_from_plain_text(file_content) + case ".json": + return _extract_text_from_json(file_content) + case ".yaml" | ".yml": + return _extract_text_from_yaml(file_content) + case ".pdf": + return _extract_text_from_pdf(file_content) + case ".doc" | ".docx": + return _extract_text_from_doc(file_content) + case ".csv": + return _extract_text_from_csv(file_content) + case ".xls" | ".xlsx": + return _extract_text_from_excel(file_content) + case ".ppt": + return _extract_text_from_ppt(file_content) + case ".pptx": + return _extract_text_from_pptx(file_content) + case ".epub": + return _extract_text_from_epub(file_content) + case ".eml": + return _extract_text_from_eml(file_content) + case ".msg": + return _extract_text_from_msg(file_content) + case _: + raise UnsupportedFileTypeError(f"Unsupported Extension Type: {file_extension}") + + +def _extract_text_from_plain_text(file_content: bytes) -> str: + try: + return file_content.decode("utf-8", "ignore") + except UnicodeDecodeError as e: + raise TextExtractionError("Failed to decode plain text file") from e + + +def _extract_text_from_json(file_content: bytes) -> str: + try: + json_data = json.loads(file_content.decode("utf-8", "ignore")) + return json.dumps(json_data, indent=2, ensure_ascii=False) + except (UnicodeDecodeError, json.JSONDecodeError) as e: + raise TextExtractionError(f"Failed to decode or parse JSON file: {e}") from e + + +def _extract_text_from_yaml(file_content: bytes) -> str: + """Extract the content from yaml file""" + try: + yaml_data = yaml.safe_load_all(file_content.decode("utf-8", "ignore")) + return cast(str, yaml.dump_all(yaml_data, allow_unicode=True, sort_keys=False)) + except (UnicodeDecodeError, yaml.YAMLError) as e: + raise TextExtractionError(f"Failed to decode or parse YAML file: {e}") from e + + +def _extract_text_from_pdf(file_content: bytes) -> str: + try: + pdf_file = io.BytesIO(file_content) + pdf_document = pypdfium2.PdfDocument(pdf_file, autoclose=True) + text = "" + for page in pdf_document: + text_page = page.get_textpage() + text += text_page.get_text_range() + text_page.close() + page.close() + return text + except Exception as e: + raise TextExtractionError(f"Failed to extract text from PDF: {str(e)}") from e + + +def _extract_text_from_doc(file_content: bytes) -> str: + """ + Extract text from a DOC/DOCX file. + For now support only paragraph and table add more if needed + """ + try: + doc_file = io.BytesIO(file_content) + doc = docx.Document(doc_file) + text = [] + + # Keep track of paragraph and table positions + content_items: list[tuple[int, str, Table | Paragraph]] = [] + + # Process paragraphs and tables + for i, paragraph in enumerate(doc.paragraphs): + if paragraph.text.strip(): + content_items.append((i, "paragraph", paragraph)) + + for i, table in enumerate(doc.tables): + content_items.append((i, "table", table)) + + # Sort content items based on their original position + content_items.sort(key=operator.itemgetter(0)) + + # Process sorted content + for _, item_type, item in content_items: + if item_type == "paragraph": + if isinstance(item, Table): + continue + text.append(item.text) + elif item_type == "table": + # Process tables + if not isinstance(item, Table): + continue + try: + # Check if any cell in the table has text + has_content = False + for row in item.rows: + if any(cell.text.strip() for cell in row.cells): + has_content = True + break + + if has_content: + cell_texts = [cell.text.replace("\n", "
") for cell in item.rows[0].cells] + markdown_table = f"| {' | '.join(cell_texts)} |\n" + markdown_table += f"| {' | '.join(['---'] * len(item.rows[0].cells))} |\n" + + for row in item.rows[1:]: + # Replace newlines with
in each cell + row_cells = [cell.text.replace("\n", "
") for cell in row.cells] + markdown_table += "| " + " | ".join(row_cells) + " |\n" + + text.append(markdown_table) + except Exception as e: + logger.warning(f"Failed to extract table from DOC/DOCX: {e}") + continue + + return "\n".join(text) + + except Exception as e: + raise TextExtractionError(f"Failed to extract text from DOC/DOCX: {str(e)}") from e + + +def _download_file_content(file: File) -> bytes: + """Download the content of a file based on its transfer method.""" + try: + if file.transfer_method == FileTransferMethod.REMOTE_URL: + if file.remote_url is None: + raise FileDownloadError("Missing URL for remote file") + response = ssrf_proxy.get(file.remote_url) + response.raise_for_status() + return cast(bytes, response.content) + else: + return cast(bytes, file_manager.download(file)) + except Exception as e: + raise FileDownloadError(f"Error downloading file: {str(e)}") from e + + +def _extract_text_from_file(file: File): + file_content = _download_file_content(file) + if file.extension: + extracted_text = _extract_text_by_file_extension(file_content=file_content, file_extension=file.extension) + elif file.mime_type: + extracted_text = _extract_text_by_mime_type(file_content=file_content, mime_type=file.mime_type) + else: + raise UnsupportedFileTypeError("Unable to determine file type: MIME type or file extension is missing") + return extracted_text + + +def _extract_text_from_csv(file_content: bytes) -> str: + try: + csv_file = io.StringIO(file_content.decode("utf-8", "ignore")) + csv_reader = csv.reader(csv_file) + rows = list(csv_reader) + + if not rows: + return "" + + # Create Markdown table + markdown_table = "| " + " | ".join(rows[0]) + " |\n" + markdown_table += "| " + " | ".join(["---"] * len(rows[0])) + " |\n" + for row in rows[1:]: + markdown_table += "| " + " | ".join(row) + " |\n" + + return markdown_table.strip() + except Exception as e: + raise TextExtractionError(f"Failed to extract text from CSV: {str(e)}") from e + + +def _extract_text_from_excel(file_content: bytes) -> str: + """Extract text from an Excel file using pandas.""" + try: + excel_file = pd.ExcelFile(io.BytesIO(file_content)) + markdown_table = "" + for sheet_name in excel_file.sheet_names: + try: + df = excel_file.parse(sheet_name=sheet_name) + df.dropna(how="all", inplace=True) + # Create Markdown table two times to separate tables with a newline + markdown_table += df.to_markdown(index=False) + "\n\n" + except Exception as e: + continue + return markdown_table + except Exception as e: + raise TextExtractionError(f"Failed to extract text from Excel file: {str(e)}") from e + + +def _extract_text_from_ppt(file_content: bytes) -> str: + from unstructured.partition.ppt import partition_ppt + + try: + with io.BytesIO(file_content) as file: + elements = partition_ppt(file=file) + return "\n".join([getattr(element, "text", "") for element in elements]) + except Exception as e: + raise TextExtractionError(f"Failed to extract text from PPT: {str(e)}") from e + + +def _extract_text_from_pptx(file_content: bytes) -> str: + from unstructured.partition.api import partition_via_api + from unstructured.partition.pptx import partition_pptx + + try: + if dify_config.UNSTRUCTURED_API_URL and dify_config.UNSTRUCTURED_API_KEY: + with tempfile.NamedTemporaryFile(suffix=".pptx", delete=False) as temp_file: + temp_file.write(file_content) + temp_file.flush() + with open(temp_file.name, "rb") as file: + elements = partition_via_api( + file=file, + metadata_filename=temp_file.name, + api_url=dify_config.UNSTRUCTURED_API_URL, + api_key=dify_config.UNSTRUCTURED_API_KEY, + ) + os.unlink(temp_file.name) + else: + with io.BytesIO(file_content) as file: + elements = partition_pptx(file=file) + return "\n".join([getattr(element, "text", "") for element in elements]) + except Exception as e: + raise TextExtractionError(f"Failed to extract text from PPTX: {str(e)}") from e + + +def _extract_text_from_epub(file_content: bytes) -> str: + from unstructured.partition.epub import partition_epub + + try: + with io.BytesIO(file_content) as file: + elements = partition_epub(file=file) + return "\n".join([str(element) for element in elements]) + except Exception as e: + raise TextExtractionError(f"Failed to extract text from EPUB: {str(e)}") from e + + +def _extract_text_from_eml(file_content: bytes) -> str: + from unstructured.partition.email import partition_email + + try: + with io.BytesIO(file_content) as file: + elements = partition_email(file=file) + return "\n".join([str(element) for element in elements]) + except Exception as e: + raise TextExtractionError(f"Failed to extract text from EML: {str(e)}") from e + + +def _extract_text_from_msg(file_content: bytes) -> str: + from unstructured.partition.msg import partition_msg + + try: + with io.BytesIO(file_content) as file: + elements = partition_msg(file=file) + return "\n".join([str(element) for element in elements]) + except Exception as e: + raise TextExtractionError(f"Failed to extract text from MSG: {str(e)}") from e diff --git a/api/core/workflow/nodes/end/__init__.py b/api/core/workflow/nodes/end/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c4c00e3ddc72d6ec502ee8f36fa7fe25c150d3f4 --- /dev/null +++ b/api/core/workflow/nodes/end/__init__.py @@ -0,0 +1,4 @@ +from .end_node import EndNode +from .entities import EndStreamParam + +__all__ = ["EndNode", "EndStreamParam"] diff --git a/api/core/workflow/nodes/end/end_node.py b/api/core/workflow/nodes/end/end_node.py new file mode 100644 index 0000000000000000000000000000000000000000..2398e4e89d59fa86ea7efb5756665e314b4d08b5 --- /dev/null +++ b/api/core/workflow/nodes/end/end_node.py @@ -0,0 +1,49 @@ +from collections.abc import Mapping, Sequence +from typing import Any + +from core.workflow.entities.node_entities import NodeRunResult +from core.workflow.nodes.base import BaseNode +from core.workflow.nodes.end.entities import EndNodeData +from core.workflow.nodes.enums import NodeType +from models.workflow import WorkflowNodeExecutionStatus + + +class EndNode(BaseNode[EndNodeData]): + _node_data_cls = EndNodeData + _node_type = NodeType.END + + def _run(self) -> NodeRunResult: + """ + Run node + :return: + """ + output_variables = self.node_data.outputs + + outputs = {} + for variable_selector in output_variables: + variable = self.graph_runtime_state.variable_pool.get(variable_selector.value_selector) + value = variable.to_object() if variable is not None else None + outputs[variable_selector.variable] = value + + return NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + inputs=outputs, + outputs=outputs, + ) + + @classmethod + def _extract_variable_selector_to_variable_mapping( + cls, + *, + graph_config: Mapping[str, Any], + node_id: str, + node_data: EndNodeData, + ) -> Mapping[str, Sequence[str]]: + """ + Extract variable selector to variable mapping + :param graph_config: graph config + :param node_id: node id + :param node_data: node data + :return: + """ + return {} diff --git a/api/core/workflow/nodes/end/end_stream_generate_router.py b/api/core/workflow/nodes/end/end_stream_generate_router.py new file mode 100644 index 0000000000000000000000000000000000000000..b3678a82b73959aa788c64ba03cb0e27ee9e884c --- /dev/null +++ b/api/core/workflow/nodes/end/end_stream_generate_router.py @@ -0,0 +1,152 @@ +from core.workflow.nodes.end.entities import EndNodeData, EndStreamParam +from core.workflow.nodes.enums import NodeType + + +class EndStreamGeneratorRouter: + @classmethod + def init( + cls, + node_id_config_mapping: dict[str, dict], + reverse_edge_mapping: dict[str, list["GraphEdge"]], # type: ignore[name-defined] + node_parallel_mapping: dict[str, str], + ) -> EndStreamParam: + """ + Get stream generate routes. + :return: + """ + # parse stream output node value selector of end nodes + end_stream_variable_selectors_mapping: dict[str, list[list[str]]] = {} + for end_node_id, node_config in node_id_config_mapping.items(): + if node_config.get("data", {}).get("type") != NodeType.END.value: + continue + + # skip end node in parallel + if end_node_id in node_parallel_mapping: + continue + + # get generate route for stream output + stream_variable_selectors = cls._extract_stream_variable_selector(node_id_config_mapping, node_config) + end_stream_variable_selectors_mapping[end_node_id] = stream_variable_selectors + + # fetch end dependencies + end_node_ids = list(end_stream_variable_selectors_mapping.keys()) + end_dependencies = cls._fetch_ends_dependencies( + end_node_ids=end_node_ids, + reverse_edge_mapping=reverse_edge_mapping, + node_id_config_mapping=node_id_config_mapping, + ) + + return EndStreamParam( + end_stream_variable_selector_mapping=end_stream_variable_selectors_mapping, + end_dependencies=end_dependencies, + ) + + @classmethod + def extract_stream_variable_selector_from_node_data( + cls, node_id_config_mapping: dict[str, dict], node_data: EndNodeData + ) -> list[list[str]]: + """ + Extract stream variable selector from node data + :param node_id_config_mapping: node id config mapping + :param node_data: node data object + :return: + """ + variable_selectors = node_data.outputs + + value_selectors = [] + for variable_selector in variable_selectors: + if not variable_selector.value_selector: + continue + + node_id = variable_selector.value_selector[0] + if node_id != "sys" and node_id in node_id_config_mapping: + node = node_id_config_mapping[node_id] + node_type = node.get("data", {}).get("type") + if ( + variable_selector.value_selector not in value_selectors + and node_type == NodeType.LLM.value + and variable_selector.value_selector[1] == "text" + ): + value_selectors.append(list(variable_selector.value_selector)) + + return value_selectors + + @classmethod + def _extract_stream_variable_selector( + cls, node_id_config_mapping: dict[str, dict], config: dict + ) -> list[list[str]]: + """ + Extract stream variable selector from node config + :param node_id_config_mapping: node id config mapping + :param config: node config + :return: + """ + node_data = EndNodeData(**config.get("data", {})) + return cls.extract_stream_variable_selector_from_node_data(node_id_config_mapping, node_data) + + @classmethod + def _fetch_ends_dependencies( + cls, + end_node_ids: list[str], + reverse_edge_mapping: dict[str, list["GraphEdge"]], # type: ignore[name-defined] + node_id_config_mapping: dict[str, dict], + ) -> dict[str, list[str]]: + """ + Fetch end dependencies + :param end_node_ids: end node ids + :param reverse_edge_mapping: reverse edge mapping + :param node_id_config_mapping: node id config mapping + :return: + """ + end_dependencies: dict[str, list[str]] = {} + for end_node_id in end_node_ids: + if end_dependencies.get(end_node_id) is None: + end_dependencies[end_node_id] = [] + + cls._recursive_fetch_end_dependencies( + current_node_id=end_node_id, + end_node_id=end_node_id, + node_id_config_mapping=node_id_config_mapping, + reverse_edge_mapping=reverse_edge_mapping, + end_dependencies=end_dependencies, + ) + + return end_dependencies + + @classmethod + def _recursive_fetch_end_dependencies( + cls, + current_node_id: str, + end_node_id: str, + node_id_config_mapping: dict[str, dict], + reverse_edge_mapping: dict[str, list["GraphEdge"]], # type: ignore[name-defined] + end_dependencies: dict[str, list[str]], + ) -> None: + """ + Recursive fetch end dependencies + :param current_node_id: current node id + :param end_node_id: end node id + :param node_id_config_mapping: node id config mapping + :param reverse_edge_mapping: reverse edge mapping + :param end_dependencies: end dependencies + :return: + """ + reverse_edges = reverse_edge_mapping.get(current_node_id, []) + for edge in reverse_edges: + source_node_id = edge.source_node_id + if source_node_id not in node_id_config_mapping: + continue + source_node_type = node_id_config_mapping[source_node_id].get("data", {}).get("type") + if source_node_type in { + NodeType.IF_ELSE.value, + NodeType.QUESTION_CLASSIFIER, + }: + end_dependencies[end_node_id].append(source_node_id) + else: + cls._recursive_fetch_end_dependencies( + current_node_id=source_node_id, + end_node_id=end_node_id, + node_id_config_mapping=node_id_config_mapping, + reverse_edge_mapping=reverse_edge_mapping, + end_dependencies=end_dependencies, + ) diff --git a/api/core/workflow/nodes/end/end_stream_processor.py b/api/core/workflow/nodes/end/end_stream_processor.py new file mode 100644 index 0000000000000000000000000000000000000000..a770eb951f6c8c9c9a294dcfce94bf0df4d002b6 --- /dev/null +++ b/api/core/workflow/nodes/end/end_stream_processor.py @@ -0,0 +1,187 @@ +import logging +from collections.abc import Generator + +from core.workflow.entities.variable_pool import VariablePool +from core.workflow.graph_engine.entities.event import ( + GraphEngineEvent, + NodeRunStartedEvent, + NodeRunStreamChunkEvent, + NodeRunSucceededEvent, +) +from core.workflow.graph_engine.entities.graph import Graph +from core.workflow.nodes.answer.base_stream_processor import StreamProcessor + +logger = logging.getLogger(__name__) + + +class EndStreamProcessor(StreamProcessor): + def __init__(self, graph: Graph, variable_pool: VariablePool) -> None: + super().__init__(graph, variable_pool) + self.end_stream_param = graph.end_stream_param + self.route_position = {} + for end_node_id, _ in self.end_stream_param.end_stream_variable_selector_mapping.items(): + self.route_position[end_node_id] = 0 + self.current_stream_chunk_generating_node_ids: dict[str, list[str]] = {} + self.has_output = False + self.output_node_ids: set[str] = set() + + def process(self, generator: Generator[GraphEngineEvent, None, None]) -> Generator[GraphEngineEvent, None, None]: + for event in generator: + if isinstance(event, NodeRunStartedEvent): + if event.route_node_state.node_id == self.graph.root_node_id and not self.rest_node_ids: + self.reset() + + yield event + elif isinstance(event, NodeRunStreamChunkEvent): + if event.in_iteration_id: + if self.has_output and event.node_id not in self.output_node_ids: + event.chunk_content = "\n" + event.chunk_content + + self.output_node_ids.add(event.node_id) + self.has_output = True + yield event + continue + + if event.route_node_state.node_id in self.current_stream_chunk_generating_node_ids: + stream_out_end_node_ids = self.current_stream_chunk_generating_node_ids[ + event.route_node_state.node_id + ] + else: + stream_out_end_node_ids = self._get_stream_out_end_node_ids(event) + self.current_stream_chunk_generating_node_ids[event.route_node_state.node_id] = ( + stream_out_end_node_ids + ) + + if stream_out_end_node_ids: + if self.has_output and event.node_id not in self.output_node_ids: + event.chunk_content = "\n" + event.chunk_content + + self.output_node_ids.add(event.node_id) + self.has_output = True + yield event + elif isinstance(event, NodeRunSucceededEvent): + yield event + if event.route_node_state.node_id in self.current_stream_chunk_generating_node_ids: + # update self.route_position after all stream event finished + for end_node_id in self.current_stream_chunk_generating_node_ids[event.route_node_state.node_id]: + self.route_position[end_node_id] += 1 + + del self.current_stream_chunk_generating_node_ids[event.route_node_state.node_id] + + # remove unreachable nodes + self._remove_unreachable_nodes(event) + + # generate stream outputs + yield from self._generate_stream_outputs_when_node_finished(event) + else: + yield event + + def reset(self) -> None: + self.route_position = {} + for end_node_id, _ in self.end_stream_param.end_stream_variable_selector_mapping.items(): + self.route_position[end_node_id] = 0 + self.rest_node_ids = self.graph.node_ids.copy() + self.current_stream_chunk_generating_node_ids = {} + + def _generate_stream_outputs_when_node_finished( + self, event: NodeRunSucceededEvent + ) -> Generator[GraphEngineEvent, None, None]: + """ + Generate stream outputs. + :param event: node run succeeded event + :return: + """ + for end_node_id, position in self.route_position.items(): + # all depends on end node id not in rest node ids + if event.route_node_state.node_id != end_node_id and ( + end_node_id not in self.rest_node_ids + or not all( + dep_id not in self.rest_node_ids for dep_id in self.end_stream_param.end_dependencies[end_node_id] + ) + ): + continue + + route_position = self.route_position[end_node_id] + + position = 0 + value_selectors = [] + for current_value_selectors in self.end_stream_param.end_stream_variable_selector_mapping[end_node_id]: + if position >= route_position: + value_selectors.append(current_value_selectors) + + position += 1 + + for value_selector in value_selectors: + if not value_selector: + continue + + value = self.variable_pool.get(value_selector) + + if value is None: + break + + text = value.markdown + + if text: + current_node_id = value_selector[0] + if self.has_output and current_node_id not in self.output_node_ids: + text = "\n" + text + + self.output_node_ids.add(current_node_id) + self.has_output = True + yield NodeRunStreamChunkEvent( + id=event.id, + node_id=event.node_id, + node_type=event.node_type, + node_data=event.node_data, + chunk_content=text, + from_variable_selector=value_selector, + route_node_state=event.route_node_state, + parallel_id=event.parallel_id, + parallel_start_node_id=event.parallel_start_node_id, + ) + + self.route_position[end_node_id] += 1 + + def _get_stream_out_end_node_ids(self, event: NodeRunStreamChunkEvent) -> list[str]: + """ + Is stream out support + :param event: queue text chunk event + :return: + """ + if not event.from_variable_selector: + return [] + + stream_output_value_selector = event.from_variable_selector + if not stream_output_value_selector: + return [] + + stream_out_end_node_ids = [] + for end_node_id, route_position in self.route_position.items(): + if end_node_id not in self.rest_node_ids: + continue + + # all depends on end node id not in rest node ids + if all(dep_id not in self.rest_node_ids for dep_id in self.end_stream_param.end_dependencies[end_node_id]): + if route_position >= len(self.end_stream_param.end_stream_variable_selector_mapping[end_node_id]): + continue + + position = 0 + value_selector = None + for current_value_selectors in self.end_stream_param.end_stream_variable_selector_mapping[end_node_id]: + if position == route_position: + value_selector = current_value_selectors + break + + position += 1 + + if not value_selector: + continue + + # check chunk node id is before current node id or equal to current node id + if value_selector != stream_output_value_selector: + continue + + stream_out_end_node_ids.append(end_node_id) + + return stream_out_end_node_ids diff --git a/api/core/workflow/nodes/end/entities.py b/api/core/workflow/nodes/end/entities.py new file mode 100644 index 0000000000000000000000000000000000000000..c16e85b0eb2a867441975da91347660804c19b77 --- /dev/null +++ b/api/core/workflow/nodes/end/entities.py @@ -0,0 +1,25 @@ +from pydantic import BaseModel, Field + +from core.workflow.entities.variable_entities import VariableSelector +from core.workflow.nodes.base import BaseNodeData + + +class EndNodeData(BaseNodeData): + """ + END Node Data. + """ + + outputs: list[VariableSelector] + + +class EndStreamParam(BaseModel): + """ + EndStreamParam entity + """ + + end_dependencies: dict[str, list[str]] = Field( + ..., description="end dependencies (end node id -> dependent node ids)" + ) + end_stream_variable_selector_mapping: dict[str, list[list[str]]] = Field( + ..., description="end stream variable selector mapping (end node id -> stream variable selectors)" + ) diff --git a/api/core/workflow/nodes/enums.py b/api/core/workflow/nodes/enums.py new file mode 100644 index 0000000000000000000000000000000000000000..7970a49aa42df47ab1bb0d7956f212a9e711e913 --- /dev/null +++ b/api/core/workflow/nodes/enums.py @@ -0,0 +1,38 @@ +from enum import StrEnum + + +class NodeType(StrEnum): + START = "start" + END = "end" + ANSWER = "answer" + LLM = "llm" + KNOWLEDGE_RETRIEVAL = "knowledge-retrieval" + IF_ELSE = "if-else" + CODE = "code" + TEMPLATE_TRANSFORM = "template-transform" + QUESTION_CLASSIFIER = "question-classifier" + HTTP_REQUEST = "http-request" + TOOL = "tool" + VARIABLE_AGGREGATOR = "variable-aggregator" + LEGACY_VARIABLE_AGGREGATOR = "variable-assigner" # TODO: Merge this into VARIABLE_AGGREGATOR in the database. + LOOP = "loop" + ITERATION = "iteration" + ITERATION_START = "iteration-start" # Fake start node for iteration. + PARAMETER_EXTRACTOR = "parameter-extractor" + VARIABLE_ASSIGNER = "assigner" + DOCUMENT_EXTRACTOR = "document-extractor" + LIST_OPERATOR = "list-operator" + + +class ErrorStrategy(StrEnum): + FAIL_BRANCH = "fail-branch" + DEFAULT_VALUE = "default-value" + + +class FailBranchSourceHandle(StrEnum): + FAILED = "fail-branch" + SUCCESS = "success-branch" + + +CONTINUE_ON_ERROR_NODE_TYPE = [NodeType.LLM, NodeType.CODE, NodeType.TOOL, NodeType.HTTP_REQUEST] +RETRY_ON_ERROR_NODE_TYPE = CONTINUE_ON_ERROR_NODE_TYPE diff --git a/api/core/workflow/nodes/event/__init__.py b/api/core/workflow/nodes/event/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..08c47d5e57387b0c5ccc85bcac2757613d86a316 --- /dev/null +++ b/api/core/workflow/nodes/event/__init__.py @@ -0,0 +1,17 @@ +from .event import ( + ModelInvokeCompletedEvent, + RunCompletedEvent, + RunRetrieverResourceEvent, + RunRetryEvent, + RunStreamChunkEvent, +) +from .types import NodeEvent + +__all__ = [ + "ModelInvokeCompletedEvent", + "NodeEvent", + "RunCompletedEvent", + "RunRetrieverResourceEvent", + "RunRetryEvent", + "RunStreamChunkEvent", +] diff --git a/api/core/workflow/nodes/event/event.py b/api/core/workflow/nodes/event/event.py new file mode 100644 index 0000000000000000000000000000000000000000..9fea3fbda3141f31fd60b6f264f54e8bacf01ce6 --- /dev/null +++ b/api/core/workflow/nodes/event/event.py @@ -0,0 +1,47 @@ +from datetime import datetime + +from pydantic import BaseModel, Field + +from core.model_runtime.entities.llm_entities import LLMUsage +from core.workflow.entities.node_entities import NodeRunResult +from models.workflow import WorkflowNodeExecutionStatus + + +class RunCompletedEvent(BaseModel): + run_result: NodeRunResult = Field(..., description="run result") + + +class RunStreamChunkEvent(BaseModel): + chunk_content: str = Field(..., description="chunk content") + from_variable_selector: list[str] = Field(..., description="from variable selector") + + +class RunRetrieverResourceEvent(BaseModel): + retriever_resources: list[dict] = Field(..., description="retriever resources") + context: str = Field(..., description="context") + + +class ModelInvokeCompletedEvent(BaseModel): + """ + Model invoke completed + """ + + text: str + usage: LLMUsage + finish_reason: str | None = None + + +class RunRetryEvent(BaseModel): + """Node Run Retry event""" + + error: str = Field(..., description="error") + retry_index: int = Field(..., description="Retry attempt number") + start_at: datetime = Field(..., description="Retry start time") + + +class SingleStepRetryEvent(NodeRunResult): + """Single step retry event""" + + status: WorkflowNodeExecutionStatus = WorkflowNodeExecutionStatus.RETRY + + elapsed_time: float = Field(..., description="elapsed time") diff --git a/api/core/workflow/nodes/event/types.py b/api/core/workflow/nodes/event/types.py new file mode 100644 index 0000000000000000000000000000000000000000..b19a91022df2e18bccf57d95d8ff3664d1bf8884 --- /dev/null +++ b/api/core/workflow/nodes/event/types.py @@ -0,0 +1,3 @@ +from .event import ModelInvokeCompletedEvent, RunCompletedEvent, RunRetrieverResourceEvent, RunStreamChunkEvent + +NodeEvent = RunCompletedEvent | RunStreamChunkEvent | RunRetrieverResourceEvent | ModelInvokeCompletedEvent diff --git a/api/core/workflow/nodes/http_request/__init__.py b/api/core/workflow/nodes/http_request/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c51c67899998327b07ac0e666d9ff518c4241280 --- /dev/null +++ b/api/core/workflow/nodes/http_request/__init__.py @@ -0,0 +1,4 @@ +from .entities import BodyData, HttpRequestNodeAuthorization, HttpRequestNodeBody, HttpRequestNodeData +from .node import HttpRequestNode + +__all__ = ["BodyData", "HttpRequestNode", "HttpRequestNodeAuthorization", "HttpRequestNodeBody", "HttpRequestNodeData"] diff --git a/api/core/workflow/nodes/http_request/entities.py b/api/core/workflow/nodes/http_request/entities.py new file mode 100644 index 0000000000000000000000000000000000000000..5764ce725ec6cec7dcfbcc1351b34611ef9fa20e --- /dev/null +++ b/api/core/workflow/nodes/http_request/entities.py @@ -0,0 +1,180 @@ +import mimetypes +from collections.abc import Sequence +from email.message import Message +from typing import Any, Literal, Optional + +import httpx +from pydantic import BaseModel, Field, ValidationInfo, field_validator + +from configs import dify_config +from core.workflow.nodes.base import BaseNodeData + + +class HttpRequestNodeAuthorizationConfig(BaseModel): + type: Literal["basic", "bearer", "custom"] + api_key: str + header: str = "" + + +class HttpRequestNodeAuthorization(BaseModel): + type: Literal["no-auth", "api-key"] + config: Optional[HttpRequestNodeAuthorizationConfig] = None + + @field_validator("config", mode="before") + @classmethod + def check_config(cls, v: HttpRequestNodeAuthorizationConfig, values: ValidationInfo): + """ + Check config, if type is no-auth, config should be None, otherwise it should be a dict. + """ + if values.data["type"] == "no-auth": + return None + else: + if not v or not isinstance(v, dict): + raise ValueError("config should be a dict") + + return v + + +class BodyData(BaseModel): + key: str = "" + type: Literal["file", "text"] + value: str = "" + file: Sequence[str] = Field(default_factory=list) + + +class HttpRequestNodeBody(BaseModel): + type: Literal["none", "form-data", "x-www-form-urlencoded", "raw-text", "json", "binary"] + data: Sequence[BodyData] = Field(default_factory=list) + + @field_validator("data", mode="before") + @classmethod + def check_data(cls, v: Any): + """For compatibility, if body is not set, return empty list.""" + if not v: + return [] + if isinstance(v, str): + return [BodyData(key="", type="text", value=v)] + return v + + +class HttpRequestNodeTimeout(BaseModel): + connect: int = dify_config.HTTP_REQUEST_MAX_CONNECT_TIMEOUT + read: int = dify_config.HTTP_REQUEST_MAX_READ_TIMEOUT + write: int = dify_config.HTTP_REQUEST_MAX_WRITE_TIMEOUT + + +class HttpRequestNodeData(BaseNodeData): + """ + Code Node Data. + """ + + method: Literal[ + "get", + "post", + "put", + "patch", + "delete", + "head", + "options", + "GET", + "POST", + "PUT", + "PATCH", + "DELETE", + "HEAD", + "OPTIONS", + ] + url: str + authorization: HttpRequestNodeAuthorization + headers: str + params: str + body: Optional[HttpRequestNodeBody] = None + timeout: Optional[HttpRequestNodeTimeout] = None + + +class Response: + headers: dict[str, str] + response: httpx.Response + + def __init__(self, response: httpx.Response): + self.response = response + self.headers = dict(response.headers) + + @property + def is_file(self): + """ + Determine if the response contains a file by checking: + 1. Content-Disposition header (RFC 6266) + 2. Content characteristics + 3. MIME type analysis + """ + content_type = self.content_type.split(";")[0].strip().lower() + content_disposition = self.response.headers.get("content-disposition", "") + + # Check if it's explicitly marked as an attachment + if content_disposition: + msg = Message() + msg["content-disposition"] = content_disposition + disp_type = msg.get_content_disposition() # Returns 'attachment', 'inline', or None + filename = msg.get_filename() # Returns filename if present, None otherwise + if disp_type == "attachment" or filename is not None: + return True + + # For application types, try to detect if it's a text-based format + if content_type.startswith("application/"): + # Common text-based application types + if any( + text_type in content_type + for text_type in ("json", "xml", "javascript", "x-www-form-urlencoded", "yaml", "graphql") + ): + return False + + # Try to detect if content is text-based by sampling first few bytes + try: + # Sample first 1024 bytes for text detection + content_sample = self.response.content[:1024] + content_sample.decode("utf-8") + # If we can decode as UTF-8 and find common text patterns, likely not a file + text_markers = (b"{", b"[", b"<", b"function", b"var ", b"const ", b"let ") + if any(marker in content_sample for marker in text_markers): + return False + except UnicodeDecodeError: + # If we can't decode as UTF-8, likely a binary file + return True + + # For other types, use MIME type analysis + main_type, _ = mimetypes.guess_type("dummy" + (mimetypes.guess_extension(content_type) or "")) + if main_type: + return main_type.split("/")[0] in ("application", "image", "audio", "video") + + # For unknown types, check if it's a media type + return any(media_type in content_type for media_type in ("image/", "audio/", "video/")) + + @property + def content_type(self) -> str: + return self.headers.get("content-type", "") + + @property + def text(self) -> str: + return self.response.text + + @property + def content(self) -> bytes: + return self.response.content + + @property + def status_code(self) -> int: + return self.response.status_code + + @property + def size(self) -> int: + return len(self.content) + + @property + def readable_size(self) -> str: + if self.size < 1024: + return f"{self.size} bytes" + elif self.size < 1024 * 1024: + return f"{(self.size / 1024):.2f} KB" + else: + return f"{(self.size / 1024 / 1024):.2f} MB" diff --git a/api/core/workflow/nodes/http_request/exc.py b/api/core/workflow/nodes/http_request/exc.py new file mode 100644 index 0000000000000000000000000000000000000000..46613c9e861bc16b58745bdc1e566d745963267b --- /dev/null +++ b/api/core/workflow/nodes/http_request/exc.py @@ -0,0 +1,26 @@ +class HttpRequestNodeError(ValueError): + """Custom error for HTTP request node.""" + + +class AuthorizationConfigError(HttpRequestNodeError): + """Raised when authorization config is missing or invalid.""" + + +class FileFetchError(HttpRequestNodeError): + """Raised when a file cannot be fetched.""" + + +class InvalidHttpMethodError(HttpRequestNodeError): + """Raised when an invalid HTTP method is used.""" + + +class ResponseSizeError(HttpRequestNodeError): + """Raised when the response size exceeds the allowed threshold.""" + + +class RequestBodyError(HttpRequestNodeError): + """Raised when the request body is invalid.""" + + +class InvalidURLError(HttpRequestNodeError): + """Raised when the URL is invalid.""" diff --git a/api/core/workflow/nodes/http_request/executor.py b/api/core/workflow/nodes/http_request/executor.py new file mode 100644 index 0000000000000000000000000000000000000000..5ed2cd6164a0ddc76ddf7b57e4eccbbcc7de319d --- /dev/null +++ b/api/core/workflow/nodes/http_request/executor.py @@ -0,0 +1,394 @@ +import json +from collections.abc import Mapping +from copy import deepcopy +from random import randint +from typing import Any, Literal +from urllib.parse import urlencode, urlparse + +import httpx + +from configs import dify_config +from core.file import file_manager +from core.helper import ssrf_proxy +from core.workflow.entities.variable_pool import VariablePool + +from .entities import ( + HttpRequestNodeAuthorization, + HttpRequestNodeData, + HttpRequestNodeTimeout, + Response, +) +from .exc import ( + AuthorizationConfigError, + FileFetchError, + HttpRequestNodeError, + InvalidHttpMethodError, + InvalidURLError, + RequestBodyError, + ResponseSizeError, +) + +BODY_TYPE_TO_CONTENT_TYPE = { + "json": "application/json", + "x-www-form-urlencoded": "application/x-www-form-urlencoded", + "form-data": "multipart/form-data", + "raw-text": "text/plain", +} + + +class Executor: + method: Literal[ + "get", + "head", + "post", + "put", + "delete", + "patch", + "options", + "GET", + "POST", + "PUT", + "PATCH", + "DELETE", + "HEAD", + "OPTIONS", + ] + url: str + params: list[tuple[str, str]] | None + content: str | bytes | None + data: Mapping[str, Any] | None + files: Mapping[str, tuple[str | None, bytes, str]] | None + json: Any + headers: dict[str, str] + auth: HttpRequestNodeAuthorization + timeout: HttpRequestNodeTimeout + max_retries: int + + boundary: str + + def __init__( + self, + *, + node_data: HttpRequestNodeData, + timeout: HttpRequestNodeTimeout, + variable_pool: VariablePool, + max_retries: int = dify_config.SSRF_DEFAULT_MAX_RETRIES, + ): + # If authorization API key is present, convert the API key using the variable pool + if node_data.authorization.type == "api-key": + if node_data.authorization.config is None: + raise AuthorizationConfigError("authorization config is required") + node_data.authorization.config.api_key = variable_pool.convert_template( + node_data.authorization.config.api_key + ).text + + self.url: str = node_data.url + self.method = node_data.method + self.auth = node_data.authorization + self.timeout = timeout + self.params = [] + self.headers = {} + self.content = None + self.files = None + self.data = None + self.json = None + self.max_retries = max_retries + + # init template + self.variable_pool = variable_pool + self.node_data = node_data + self._initialize() + + def _initialize(self): + self._init_url() + self._init_params() + self._init_headers() + self._init_body() + + def _init_url(self): + self.url = self.variable_pool.convert_template(self.node_data.url).text + + # check if url is a valid URL + if not self.url: + raise InvalidURLError("url is required") + if not self.url.startswith(("http://", "https://")): + raise InvalidURLError("url should start with http:// or https://") + + def _init_params(self): + """ + Almost same as _init_headers(), difference: + 1. response a list tuple to support same key, like 'aa=1&aa=2' + 2. param value may have '\n', we need to splitlines then extract the variable value. + """ + result = [] + for line in self.node_data.params.splitlines(): + if not (line := line.strip()): + continue + + key, *value = line.split(":", 1) + if not (key := key.strip()): + continue + + value_str = value[0].strip() if value else "" + result.append( + (self.variable_pool.convert_template(key).text, self.variable_pool.convert_template(value_str).text) + ) + + self.params = result + + def _init_headers(self): + """ + Convert the header string of frontend to a dictionary. + + Each line in the header string represents a key-value pair. + Keys and values are separated by ':'. + Empty values are allowed. + + Examples: + 'aa:bb\n cc:dd' -> {'aa': 'bb', 'cc': 'dd'} + 'aa:\n cc:dd\n' -> {'aa': '', 'cc': 'dd'} + 'aa\n cc : dd' -> {'aa': '', 'cc': 'dd'} + + """ + headers = self.variable_pool.convert_template(self.node_data.headers).text + self.headers = { + key.strip(): (value[0].strip() if value else "") + for line in headers.splitlines() + if line.strip() + for key, *value in [line.split(":", 1)] + } + + def _init_body(self): + body = self.node_data.body + if body is not None: + data = body.data + match body.type: + case "none": + self.content = "" + case "raw-text": + if len(data) != 1: + raise RequestBodyError("raw-text body type should have exactly one item") + self.content = self.variable_pool.convert_template(data[0].value).text + case "json": + if len(data) != 1: + raise RequestBodyError("json body type should have exactly one item") + json_string = self.variable_pool.convert_template(data[0].value).text + try: + json_object = json.loads(json_string, strict=False) + except json.JSONDecodeError as e: + raise RequestBodyError(f"Failed to parse JSON: {json_string}") from e + self.json = json_object + # self.json = self._parse_object_contains_variables(json_object) + case "binary": + if len(data) != 1: + raise RequestBodyError("binary body type should have exactly one item") + file_selector = data[0].file + file_variable = self.variable_pool.get_file(file_selector) + if file_variable is None: + raise FileFetchError(f"cannot fetch file with selector {file_selector}") + file = file_variable.value + self.content = file_manager.download(file) + case "x-www-form-urlencoded": + form_data = { + self.variable_pool.convert_template(item.key).text: self.variable_pool.convert_template( + item.value + ).text + for item in data + } + self.data = form_data + case "form-data": + form_data = { + self.variable_pool.convert_template(item.key).text: self.variable_pool.convert_template( + item.value + ).text + for item in filter(lambda item: item.type == "text", data) + } + file_selectors = { + self.variable_pool.convert_template(item.key).text: item.file + for item in filter(lambda item: item.type == "file", data) + } + files: dict[str, Any] = {} + files = {k: self.variable_pool.get_file(selector) for k, selector in file_selectors.items()} + files = {k: v for k, v in files.items() if v is not None} + files = {k: variable.value for k, variable in files.items() if variable is not None} + files = { + k: (v.filename, file_manager.download(v), v.mime_type or "application/octet-stream") + for k, v in files.items() + if v.related_id is not None + } + self.data = form_data + self.files = files or None + + def _assembling_headers(self) -> dict[str, Any]: + authorization = deepcopy(self.auth) + headers = deepcopy(self.headers) or {} + if self.auth.type == "api-key": + if self.auth.config is None: + raise AuthorizationConfigError("self.authorization config is required") + if authorization.config is None: + raise AuthorizationConfigError("authorization config is required") + + if self.auth.config.api_key is None: + raise AuthorizationConfigError("api_key is required") + + if not authorization.config.header: + authorization.config.header = "Authorization" + + if self.auth.config.type == "bearer": + headers[authorization.config.header] = f"Bearer {authorization.config.api_key}" + elif self.auth.config.type == "basic": + headers[authorization.config.header] = f"Basic {authorization.config.api_key}" + elif self.auth.config.type == "custom": + headers[authorization.config.header] = authorization.config.api_key or "" + + return headers + + def _validate_and_parse_response(self, response: httpx.Response) -> Response: + executor_response = Response(response) + + threshold_size = ( + dify_config.HTTP_REQUEST_NODE_MAX_BINARY_SIZE + if executor_response.is_file + else dify_config.HTTP_REQUEST_NODE_MAX_TEXT_SIZE + ) + if executor_response.size > threshold_size: + raise ResponseSizeError( + f"{'File' if executor_response.is_file else 'Text'} size is too large," + f" max size is {threshold_size / 1024 / 1024:.2f} MB," + f" but current size is {executor_response.readable_size}." + ) + + return executor_response + + def _do_http_request(self, headers: dict[str, Any]) -> httpx.Response: + """ + do http request depending on api bundle + """ + if self.method not in { + "get", + "head", + "post", + "put", + "delete", + "patch", + "options", + "GET", + "POST", + "PUT", + "PATCH", + "DELETE", + "HEAD", + "OPTIONS", + }: + raise InvalidHttpMethodError(f"Invalid http method {self.method}") + + request_args = { + "url": self.url, + "data": self.data, + "files": self.files, + "json": self.json, + "content": self.content, + "headers": headers, + "params": self.params, + "timeout": (self.timeout.connect, self.timeout.read, self.timeout.write), + "follow_redirects": True, + "max_retries": self.max_retries, + } + # request_args = {k: v for k, v in request_args.items() if v is not None} + try: + response = getattr(ssrf_proxy, self.method.lower())(**request_args) + except (ssrf_proxy.MaxRetriesExceededError, httpx.RequestError) as e: + raise HttpRequestNodeError(str(e)) + # FIXME: fix type ignore, this maybe httpx type issue + return response # type: ignore + + def invoke(self) -> Response: + # assemble headers + headers = self._assembling_headers() + # do http request + response = self._do_http_request(headers) + # validate response + return self._validate_and_parse_response(response) + + def to_log(self): + url_parts = urlparse(self.url) + path = url_parts.path or "/" + + # Add query parameters + if self.params: + query_string = urlencode(self.params) + path += f"?{query_string}" + elif url_parts.query: + path += f"?{url_parts.query}" + + raw = f"{self.method.upper()} {path} HTTP/1.1\r\n" + raw += f"Host: {url_parts.netloc}\r\n" + + headers = self._assembling_headers() + body = self.node_data.body + boundary = f"----WebKitFormBoundary{_generate_random_string(16)}" + if body: + if "content-type" not in (k.lower() for k in self.headers) and body.type in BODY_TYPE_TO_CONTENT_TYPE: + headers["Content-Type"] = BODY_TYPE_TO_CONTENT_TYPE[body.type] + if body.type == "form-data": + headers["Content-Type"] = f"multipart/form-data; boundary={boundary}" + for k, v in headers.items(): + if self.auth.type == "api-key": + authorization_header = "Authorization" + if self.auth.config and self.auth.config.header: + authorization_header = self.auth.config.header + if k.lower() == authorization_header.lower(): + raw += f"{k}: {'*' * len(v)}\r\n" + continue + raw += f"{k}: {v}\r\n" + + body_string = "" + if self.files: + for k, v in self.files.items(): + body_string += f"--{boundary}\r\n" + body_string += f'Content-Disposition: form-data; name="{k}"\r\n\r\n' + body_string += f"{v[1]}\r\n" + body_string += f"--{boundary}--\r\n" + elif self.node_data.body: + if self.content: + if isinstance(self.content, str): + body_string = self.content + elif isinstance(self.content, bytes): + body_string = self.content.decode("utf-8", errors="replace") + elif self.data and self.node_data.body.type == "x-www-form-urlencoded": + body_string = urlencode(self.data) + elif self.data and self.node_data.body.type == "form-data": + for key, value in self.data.items(): + body_string += f"--{boundary}\r\n" + body_string += f'Content-Disposition: form-data; name="{key}"\r\n\r\n' + body_string += f"{value}\r\n" + body_string += f"--{boundary}--\r\n" + elif self.json: + body_string = json.dumps(self.json) + elif self.node_data.body.type == "raw-text": + if len(self.node_data.body.data) != 1: + raise RequestBodyError("raw-text body type should have exactly one item") + body_string = self.node_data.body.data[0].value + if body_string: + raw += f"Content-Length: {len(body_string)}\r\n" + raw += "\r\n" # Empty line between headers and body + raw += body_string + + return raw + + +def _generate_random_string(n: int) -> str: + """ + Generate a random string of lowercase ASCII letters. + + Args: + n (int): The length of the random string to generate. + + Returns: + str: A random string of lowercase ASCII letters with length n. + + Example: + >>> _generate_random_string(5) + 'abcde' + """ + return "".join([chr(randint(97, 122)) for _ in range(n)]) diff --git a/api/core/workflow/nodes/http_request/node.py b/api/core/workflow/nodes/http_request/node.py new file mode 100644 index 0000000000000000000000000000000000000000..861119f26cb08882ed82e023bc3f519fe20e3fd3 --- /dev/null +++ b/api/core/workflow/nodes/http_request/node.py @@ -0,0 +1,200 @@ +import logging +import mimetypes +from collections.abc import Mapping, Sequence +from typing import Any, Optional + +from configs import dify_config +from core.file import File, FileTransferMethod +from core.tools.tool_file_manager import ToolFileManager +from core.workflow.entities.node_entities import NodeRunResult +from core.workflow.entities.variable_entities import VariableSelector +from core.workflow.nodes.base import BaseNode +from core.workflow.nodes.enums import NodeType +from core.workflow.nodes.http_request.executor import Executor +from core.workflow.utils import variable_template_parser +from factories import file_factory +from models.workflow import WorkflowNodeExecutionStatus + +from .entities import ( + HttpRequestNodeData, + HttpRequestNodeTimeout, + Response, +) +from .exc import HttpRequestNodeError, RequestBodyError + +HTTP_REQUEST_DEFAULT_TIMEOUT = HttpRequestNodeTimeout( + connect=dify_config.HTTP_REQUEST_MAX_CONNECT_TIMEOUT, + read=dify_config.HTTP_REQUEST_MAX_READ_TIMEOUT, + write=dify_config.HTTP_REQUEST_MAX_WRITE_TIMEOUT, +) + +logger = logging.getLogger(__name__) + + +class HttpRequestNode(BaseNode[HttpRequestNodeData]): + _node_data_cls = HttpRequestNodeData + _node_type = NodeType.HTTP_REQUEST + + @classmethod + def get_default_config(cls, filters: Optional[dict[str, Any]] = None) -> dict: + return { + "type": "http-request", + "config": { + "method": "get", + "authorization": { + "type": "no-auth", + }, + "body": {"type": "none"}, + "timeout": { + **HTTP_REQUEST_DEFAULT_TIMEOUT.model_dump(), + "max_connect_timeout": dify_config.HTTP_REQUEST_MAX_CONNECT_TIMEOUT, + "max_read_timeout": dify_config.HTTP_REQUEST_MAX_READ_TIMEOUT, + "max_write_timeout": dify_config.HTTP_REQUEST_MAX_WRITE_TIMEOUT, + }, + }, + "retry_config": { + "max_retries": dify_config.SSRF_DEFAULT_MAX_RETRIES, + "retry_interval": 0.5 * (2**2), + "retry_enabled": True, + }, + } + + def _run(self) -> NodeRunResult: + process_data = {} + try: + http_executor = Executor( + node_data=self.node_data, + timeout=self._get_request_timeout(self.node_data), + variable_pool=self.graph_runtime_state.variable_pool, + max_retries=0, + ) + process_data["request"] = http_executor.to_log() + + response = http_executor.invoke() + files = self.extract_files(url=http_executor.url, response=response) + if not response.response.is_success and (self.should_continue_on_error or self.should_retry): + return NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + outputs={ + "status_code": response.status_code, + "body": response.text if not files else "", + "headers": response.headers, + "files": files, + }, + process_data={ + "request": http_executor.to_log(), + }, + error=f"Request failed with status code {response.status_code}", + error_type="HTTPResponseCodeError", + ) + return NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + outputs={ + "status_code": response.status_code, + "body": response.text if not files else "", + "headers": response.headers, + "files": files, + }, + process_data={ + "request": http_executor.to_log(), + }, + ) + except HttpRequestNodeError as e: + logger.warning(f"http request node {self.node_id} failed to run: {e}") + return NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + error=str(e), + process_data=process_data, + error_type=type(e).__name__, + ) + + @staticmethod + def _get_request_timeout(node_data: HttpRequestNodeData) -> HttpRequestNodeTimeout: + timeout = node_data.timeout + if timeout is None: + return HTTP_REQUEST_DEFAULT_TIMEOUT + + timeout.connect = timeout.connect or HTTP_REQUEST_DEFAULT_TIMEOUT.connect + timeout.read = timeout.read or HTTP_REQUEST_DEFAULT_TIMEOUT.read + timeout.write = timeout.write or HTTP_REQUEST_DEFAULT_TIMEOUT.write + return timeout + + @classmethod + def _extract_variable_selector_to_variable_mapping( + cls, + *, + graph_config: Mapping[str, Any], + node_id: str, + node_data: HttpRequestNodeData, + ) -> Mapping[str, Sequence[str]]: + selectors: list[VariableSelector] = [] + selectors += variable_template_parser.extract_selectors_from_template(node_data.url) + selectors += variable_template_parser.extract_selectors_from_template(node_data.headers) + selectors += variable_template_parser.extract_selectors_from_template(node_data.params) + if node_data.body: + body_type = node_data.body.type + data = node_data.body.data + match body_type: + case "binary": + if len(data) != 1: + raise RequestBodyError("invalid body data, should have only one item") + selector = data[0].file + selectors.append(VariableSelector(variable="#" + ".".join(selector) + "#", value_selector=selector)) + case "json" | "raw-text": + if len(data) != 1: + raise RequestBodyError("invalid body data, should have only one item") + selectors += variable_template_parser.extract_selectors_from_template(data[0].key) + selectors += variable_template_parser.extract_selectors_from_template(data[0].value) + case "x-www-form-urlencoded": + for item in data: + selectors += variable_template_parser.extract_selectors_from_template(item.key) + selectors += variable_template_parser.extract_selectors_from_template(item.value) + case "form-data": + for item in data: + selectors += variable_template_parser.extract_selectors_from_template(item.key) + if item.type == "text": + selectors += variable_template_parser.extract_selectors_from_template(item.value) + elif item.type == "file": + selectors.append( + VariableSelector(variable="#" + ".".join(item.file) + "#", value_selector=item.file) + ) + + mapping = {} + for selector_iter in selectors: + mapping[node_id + "." + selector_iter.variable] = selector_iter.value_selector + + return mapping + + def extract_files(self, url: str, response: Response) -> list[File]: + """ + Extract files from response by checking both Content-Type header and URL + """ + files = [] + is_file = response.is_file + content_type = response.content_type + content = response.content + + if is_file: + # Guess file extension from URL or Content-Type header + filename = url.split("?")[0].split("/")[-1] or "" + mime_type = content_type or mimetypes.guess_type(filename)[0] or "application/octet-stream" + + tool_file = ToolFileManager.create_file_by_raw( + user_id=self.user_id, + tenant_id=self.tenant_id, + conversation_id=None, + file_binary=content, + mimetype=mime_type, + ) + + mapping = { + "tool_file_id": tool_file.id, + "transfer_method": FileTransferMethod.TOOL_FILE.value, + } + file = file_factory.build_from_mapping( + mapping=mapping, + tenant_id=self.tenant_id, + ) + files.append(file) + + return files diff --git a/api/core/workflow/nodes/if_else/__init__.py b/api/core/workflow/nodes/if_else/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..afa0e8112c5b17e6fa956f4c054c29e5251381bf --- /dev/null +++ b/api/core/workflow/nodes/if_else/__init__.py @@ -0,0 +1,3 @@ +from .if_else_node import IfElseNode + +__all__ = ["IfElseNode"] diff --git a/api/core/workflow/nodes/if_else/entities.py b/api/core/workflow/nodes/if_else/entities.py new file mode 100644 index 0000000000000000000000000000000000000000..23f5d2cc317f78b623ae148489b7411b42f22313 --- /dev/null +++ b/api/core/workflow/nodes/if_else/entities.py @@ -0,0 +1,26 @@ +from typing import Literal, Optional + +from pydantic import BaseModel, Field + +from core.workflow.nodes.base import BaseNodeData +from core.workflow.utils.condition.entities import Condition + + +class IfElseNodeData(BaseNodeData): + """ + Answer Node Data. + """ + + class Case(BaseModel): + """ + Case entity representing a single logical condition group + """ + + case_id: str + logical_operator: Literal["and", "or"] + conditions: list[Condition] + + logical_operator: Optional[Literal["and", "or"]] = "and" + conditions: Optional[list[Condition]] = Field(default=None, deprecated=True) + + cases: Optional[list[Case]] = None diff --git a/api/core/workflow/nodes/if_else/if_else_node.py b/api/core/workflow/nodes/if_else/if_else_node.py new file mode 100644 index 0000000000000000000000000000000000000000..a1dc0f0664a32de246d2c2cd941ee60c9cb37d83 --- /dev/null +++ b/api/core/workflow/nodes/if_else/if_else_node.py @@ -0,0 +1,121 @@ +from collections.abc import Mapping, Sequence +from typing import Any, Literal + +from typing_extensions import deprecated + +from core.workflow.entities.node_entities import NodeRunResult +from core.workflow.entities.variable_pool import VariablePool +from core.workflow.nodes.base import BaseNode +from core.workflow.nodes.enums import NodeType +from core.workflow.nodes.if_else.entities import IfElseNodeData +from core.workflow.utils.condition.entities import Condition +from core.workflow.utils.condition.processor import ConditionProcessor +from models.workflow import WorkflowNodeExecutionStatus + + +class IfElseNode(BaseNode[IfElseNodeData]): + _node_data_cls = IfElseNodeData + _node_type = NodeType.IF_ELSE + + def _run(self) -> NodeRunResult: + """ + Run node + :return: + """ + node_inputs: dict[str, list] = {"conditions": []} + + process_data: dict[str, list] = {"condition_results": []} + + input_conditions = [] + final_result = False + selected_case_id = None + condition_processor = ConditionProcessor() + try: + # Check if the new cases structure is used + if self.node_data.cases: + for case in self.node_data.cases: + input_conditions, group_result, final_result = condition_processor.process_conditions( + variable_pool=self.graph_runtime_state.variable_pool, + conditions=case.conditions, + operator=case.logical_operator, + ) + + process_data["condition_results"].append( + { + "group": case.model_dump(), + "results": group_result, + "final_result": final_result, + } + ) + + # Break if a case passes (logical short-circuit) + if final_result: + selected_case_id = case.case_id # Capture the ID of the passing case + break + + else: + # TODO: Update database then remove this + # Fallback to old structure if cases are not defined + input_conditions, group_result, final_result = _should_not_use_old_function( + condition_processor=condition_processor, + variable_pool=self.graph_runtime_state.variable_pool, + conditions=self.node_data.conditions or [], + operator=self.node_data.logical_operator or "and", + ) + + selected_case_id = "true" if final_result else "false" + + process_data["condition_results"].append( + {"group": "default", "results": group_result, "final_result": final_result} + ) + + node_inputs["conditions"] = input_conditions + + except Exception as e: + return NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, inputs=node_inputs, process_data=process_data, error=str(e) + ) + + outputs = {"result": final_result, "selected_case_id": selected_case_id} + + data = NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + inputs=node_inputs, + process_data=process_data, + edge_source_handle=selected_case_id or "false", # Use case ID or 'default' + outputs=outputs, + ) + + return data + + @classmethod + def _extract_variable_selector_to_variable_mapping( + cls, + *, + graph_config: Mapping[str, Any], + node_id: str, + node_data: IfElseNodeData, + ) -> Mapping[str, Sequence[str]]: + """ + Extract variable selector to variable mapping + :param graph_config: graph config + :param node_id: node id + :param node_data: node data + :return: + """ + return {} + + +@deprecated("This function is deprecated. You should use the new cases structure.") +def _should_not_use_old_function( + *, + condition_processor: ConditionProcessor, + variable_pool: VariablePool, + conditions: list[Condition], + operator: Literal["and", "or"], +): + return condition_processor.process_conditions( + variable_pool=variable_pool, + conditions=conditions, + operator=operator, + ) diff --git a/api/core/workflow/nodes/iteration/__init__.py b/api/core/workflow/nodes/iteration/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5bb87aaffa92b43614fa8c990972298f12d7edfa --- /dev/null +++ b/api/core/workflow/nodes/iteration/__init__.py @@ -0,0 +1,5 @@ +from .entities import IterationNodeData +from .iteration_node import IterationNode +from .iteration_start_node import IterationStartNode + +__all__ = ["IterationNode", "IterationNodeData", "IterationStartNode"] diff --git a/api/core/workflow/nodes/iteration/entities.py b/api/core/workflow/nodes/iteration/entities.py new file mode 100644 index 0000000000000000000000000000000000000000..7a489dd725d65359a925021e5580ecf59043cdc5 --- /dev/null +++ b/api/core/workflow/nodes/iteration/entities.py @@ -0,0 +1,63 @@ +from enum import StrEnum +from typing import Any, Optional + +from pydantic import Field + +from core.workflow.nodes.base import BaseIterationNodeData, BaseIterationState, BaseNodeData + + +class ErrorHandleMode(StrEnum): + TERMINATED = "terminated" + CONTINUE_ON_ERROR = "continue-on-error" + REMOVE_ABNORMAL_OUTPUT = "remove-abnormal-output" + + +class IterationNodeData(BaseIterationNodeData): + """ + Iteration Node Data. + """ + + parent_loop_id: Optional[str] = None # redundant field, not used currently + iterator_selector: list[str] # variable selector + output_selector: list[str] # output selector + is_parallel: bool = False # open the parallel mode or not + parallel_nums: int = 10 # the numbers of parallel + error_handle_mode: ErrorHandleMode = ErrorHandleMode.TERMINATED # how to handle the error + + +class IterationStartNodeData(BaseNodeData): + """ + Iteration Start Node Data. + """ + + pass + + +class IterationState(BaseIterationState): + """ + Iteration State. + """ + + outputs: list[Any] = Field(default_factory=list) + current_output: Optional[Any] = None + + class MetaData(BaseIterationState.MetaData): + """ + Data. + """ + + iterator_length: int + + def get_last_output(self) -> Optional[Any]: + """ + Get last output. + """ + if self.outputs: + return self.outputs[-1] + return None + + def get_current_output(self) -> Optional[Any]: + """ + Get current output. + """ + return self.current_output diff --git a/api/core/workflow/nodes/iteration/exc.py b/api/core/workflow/nodes/iteration/exc.py new file mode 100644 index 0000000000000000000000000000000000000000..d9947e09bc10c844f71cffdc85d26573b22e8f70 --- /dev/null +++ b/api/core/workflow/nodes/iteration/exc.py @@ -0,0 +1,22 @@ +class IterationNodeError(ValueError): + """Base class for iteration node errors.""" + + +class IteratorVariableNotFoundError(IterationNodeError): + """Raised when the iterator variable is not found.""" + + +class InvalidIteratorValueError(IterationNodeError): + """Raised when the iterator value is invalid.""" + + +class StartNodeIdNotFoundError(IterationNodeError): + """Raised when the start node ID is not found.""" + + +class IterationGraphNotFoundError(IterationNodeError): + """Raised when the iteration graph is not found.""" + + +class IterationIndexNotFoundError(IterationNodeError): + """Raised when the iteration index is not found.""" diff --git a/api/core/workflow/nodes/iteration/iteration_node.py b/api/core/workflow/nodes/iteration/iteration_node.py new file mode 100644 index 0000000000000000000000000000000000000000..f1289558fffa82606a41d764c9322facf84a7813 --- /dev/null +++ b/api/core/workflow/nodes/iteration/iteration_node.py @@ -0,0 +1,603 @@ +import logging +import uuid +from collections.abc import Generator, Mapping, Sequence +from concurrent.futures import Future, wait +from datetime import UTC, datetime +from queue import Empty, Queue +from typing import TYPE_CHECKING, Any, Optional, cast + +from flask import Flask, current_app + +from configs import dify_config +from core.variables import ArrayVariable, IntegerVariable, NoneVariable +from core.workflow.entities.node_entities import ( + NodeRunMetadataKey, + NodeRunResult, +) +from core.workflow.entities.variable_pool import VariablePool +from core.workflow.graph_engine.entities.event import ( + BaseGraphEvent, + BaseNodeEvent, + BaseParallelBranchEvent, + GraphRunFailedEvent, + InNodeEvent, + IterationRunFailedEvent, + IterationRunNextEvent, + IterationRunStartedEvent, + IterationRunSucceededEvent, + NodeInIterationFailedEvent, + NodeRunFailedEvent, + NodeRunStartedEvent, + NodeRunStreamChunkEvent, + NodeRunSucceededEvent, +) +from core.workflow.graph_engine.entities.graph import Graph +from core.workflow.nodes.base import BaseNode +from core.workflow.nodes.enums import NodeType +from core.workflow.nodes.event import NodeEvent, RunCompletedEvent +from core.workflow.nodes.iteration.entities import ErrorHandleMode, IterationNodeData +from models.workflow import WorkflowNodeExecutionStatus + +from .exc import ( + InvalidIteratorValueError, + IterationGraphNotFoundError, + IterationIndexNotFoundError, + IterationNodeError, + IteratorVariableNotFoundError, + StartNodeIdNotFoundError, +) + +if TYPE_CHECKING: + from core.workflow.graph_engine.graph_engine import GraphEngine +logger = logging.getLogger(__name__) + + +class IterationNode(BaseNode[IterationNodeData]): + """ + Iteration Node. + """ + + _node_data_cls = IterationNodeData + _node_type = NodeType.ITERATION + + @classmethod + def get_default_config(cls, filters: Optional[dict] = None) -> dict: + return { + "type": "iteration", + "config": { + "is_parallel": False, + "parallel_nums": 10, + "error_handle_mode": ErrorHandleMode.TERMINATED.value, + }, + } + + def _run(self) -> Generator[NodeEvent | InNodeEvent, None, None]: + """ + Run the node. + """ + variable = self.graph_runtime_state.variable_pool.get(self.node_data.iterator_selector) + + if not variable: + raise IteratorVariableNotFoundError(f"iterator variable {self.node_data.iterator_selector} not found") + + if not isinstance(variable, ArrayVariable) and not isinstance(variable, NoneVariable): + raise InvalidIteratorValueError(f"invalid iterator value: {variable}, please provide a list.") + + if isinstance(variable, NoneVariable) or len(variable.value) == 0: + yield RunCompletedEvent( + run_result=NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + outputs={"output": []}, + ) + ) + return + + iterator_list_value = variable.to_object() + + if not isinstance(iterator_list_value, list): + raise InvalidIteratorValueError(f"Invalid iterator value: {iterator_list_value}, please provide a list.") + + inputs = {"iterator_selector": iterator_list_value} + + graph_config = self.graph_config + + if not self.node_data.start_node_id: + raise StartNodeIdNotFoundError(f"field start_node_id in iteration {self.node_id} not found") + + root_node_id = self.node_data.start_node_id + + # init graph + iteration_graph = Graph.init(graph_config=graph_config, root_node_id=root_node_id) + + if not iteration_graph: + raise IterationGraphNotFoundError("iteration graph not found") + + variable_pool = self.graph_runtime_state.variable_pool + + # append iteration variable (item, index) to variable pool + variable_pool.add([self.node_id, "index"], 0) + variable_pool.add([self.node_id, "item"], iterator_list_value[0]) + + # init graph engine + from core.workflow.graph_engine.graph_engine import GraphEngine, GraphEngineThreadPool + + graph_engine = GraphEngine( + tenant_id=self.tenant_id, + app_id=self.app_id, + workflow_type=self.workflow_type, + workflow_id=self.workflow_id, + user_id=self.user_id, + user_from=self.user_from, + invoke_from=self.invoke_from, + call_depth=self.workflow_call_depth, + graph=iteration_graph, + graph_config=graph_config, + variable_pool=variable_pool, + max_execution_steps=dify_config.WORKFLOW_MAX_EXECUTION_STEPS, + max_execution_time=dify_config.WORKFLOW_MAX_EXECUTION_TIME, + thread_pool_id=self.thread_pool_id, + ) + + start_at = datetime.now(UTC).replace(tzinfo=None) + + yield IterationRunStartedEvent( + iteration_id=self.id, + iteration_node_id=self.node_id, + iteration_node_type=self.node_type, + iteration_node_data=self.node_data, + start_at=start_at, + inputs=inputs, + metadata={"iterator_length": len(iterator_list_value)}, + predecessor_node_id=self.previous_node_id, + ) + + yield IterationRunNextEvent( + iteration_id=self.id, + iteration_node_id=self.node_id, + iteration_node_type=self.node_type, + iteration_node_data=self.node_data, + index=0, + pre_iteration_output=None, + duration=None, + ) + iter_run_map: dict[str, float] = {} + outputs: list[Any] = [None] * len(iterator_list_value) + try: + if self.node_data.is_parallel: + futures: list[Future] = [] + q: Queue = Queue() + thread_pool = GraphEngineThreadPool( + max_workers=self.node_data.parallel_nums, max_submit_count=dify_config.MAX_SUBMIT_COUNT + ) + for index, item in enumerate(iterator_list_value): + future: Future = thread_pool.submit( + self._run_single_iter_parallel, + flask_app=current_app._get_current_object(), # type: ignore + q=q, + iterator_list_value=iterator_list_value, + inputs=inputs, + outputs=outputs, + start_at=start_at, + graph_engine=graph_engine, + iteration_graph=iteration_graph, + index=index, + item=item, + iter_run_map=iter_run_map, + ) + future.add_done_callback(thread_pool.task_done_callback) + futures.append(future) + succeeded_count = 0 + while True: + try: + event = q.get(timeout=1) + if event is None: + break + if isinstance(event, IterationRunNextEvent): + succeeded_count += 1 + if succeeded_count == len(futures): + q.put(None) + yield event + if isinstance(event, RunCompletedEvent): + q.put(None) + for f in futures: + if not f.done(): + f.cancel() + yield event + if isinstance(event, IterationRunFailedEvent): + q.put(None) + yield event + except Empty: + continue + + # wait all threads + wait(futures) + else: + for _ in range(len(iterator_list_value)): + yield from self._run_single_iter( + iterator_list_value=iterator_list_value, + variable_pool=variable_pool, + inputs=inputs, + outputs=outputs, + start_at=start_at, + graph_engine=graph_engine, + iteration_graph=iteration_graph, + iter_run_map=iter_run_map, + ) + if self.node_data.error_handle_mode == ErrorHandleMode.REMOVE_ABNORMAL_OUTPUT: + outputs = [output for output in outputs if output is not None] + + # Flatten the list of lists + if isinstance(outputs, list) and all(isinstance(output, list) for output in outputs): + outputs = [item for sublist in outputs for item in sublist] + + yield IterationRunSucceededEvent( + iteration_id=self.id, + iteration_node_id=self.node_id, + iteration_node_type=self.node_type, + iteration_node_data=self.node_data, + start_at=start_at, + inputs=inputs, + outputs={"output": outputs}, + steps=len(iterator_list_value), + metadata={"total_tokens": graph_engine.graph_runtime_state.total_tokens}, + ) + + yield RunCompletedEvent( + run_result=NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + outputs={"output": outputs}, + metadata={ + NodeRunMetadataKey.ITERATION_DURATION_MAP: iter_run_map, + NodeRunMetadataKey.TOTAL_TOKENS: graph_engine.graph_runtime_state.total_tokens, + }, + ) + ) + except IterationNodeError as e: + # iteration run failed + logger.warning("Iteration run failed") + yield IterationRunFailedEvent( + iteration_id=self.id, + iteration_node_id=self.node_id, + iteration_node_type=self.node_type, + iteration_node_data=self.node_data, + start_at=start_at, + inputs=inputs, + outputs={"output": outputs}, + steps=len(iterator_list_value), + metadata={"total_tokens": graph_engine.graph_runtime_state.total_tokens}, + error=str(e), + ) + + yield RunCompletedEvent( + run_result=NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + error=str(e), + ) + ) + finally: + # remove iteration variable (item, index) from variable pool after iteration run completed + variable_pool.remove([self.node_id, "index"]) + variable_pool.remove([self.node_id, "item"]) + + @classmethod + def _extract_variable_selector_to_variable_mapping( + cls, + *, + graph_config: Mapping[str, Any], + node_id: str, + node_data: IterationNodeData, + ) -> Mapping[str, Sequence[str]]: + """ + Extract variable selector to variable mapping + :param graph_config: graph config + :param node_id: node id + :param node_data: node data + :return: + """ + variable_mapping: dict[str, Sequence[str]] = { + f"{node_id}.input_selector": node_data.iterator_selector, + } + + # init graph + iteration_graph = Graph.init(graph_config=graph_config, root_node_id=node_data.start_node_id) + + if not iteration_graph: + raise IterationGraphNotFoundError("iteration graph not found") + + for sub_node_id, sub_node_config in iteration_graph.node_id_config_mapping.items(): + if sub_node_config.get("data", {}).get("iteration_id") != node_id: + continue + + # variable selector to variable mapping + try: + # Get node class + from core.workflow.nodes.node_mapping import NODE_TYPE_CLASSES_MAPPING + + node_type = NodeType(sub_node_config.get("data", {}).get("type")) + if node_type not in NODE_TYPE_CLASSES_MAPPING: + continue + node_version = sub_node_config.get("data", {}).get("version", "1") + node_cls = NODE_TYPE_CLASSES_MAPPING[node_type][node_version] + + sub_node_variable_mapping = node_cls.extract_variable_selector_to_variable_mapping( + graph_config=graph_config, config=sub_node_config + ) + sub_node_variable_mapping = cast(dict[str, Sequence[str]], sub_node_variable_mapping) + except NotImplementedError: + sub_node_variable_mapping = {} + + # remove iteration variables + sub_node_variable_mapping = { + sub_node_id + "." + key: value + for key, value in sub_node_variable_mapping.items() + if value[0] != node_id + } + + variable_mapping.update(sub_node_variable_mapping) + + # remove variable out from iteration + variable_mapping = { + key: value for key, value in variable_mapping.items() if value[0] not in iteration_graph.node_ids + } + + return variable_mapping + + def _handle_event_metadata( + self, + *, + event: BaseNodeEvent | InNodeEvent, + iter_run_index: int, + parallel_mode_run_id: str | None, + ) -> NodeRunStartedEvent | BaseNodeEvent | InNodeEvent: + """ + add iteration metadata to event. + """ + if not isinstance(event, BaseNodeEvent): + return event + if self.node_data.is_parallel and isinstance(event, NodeRunStartedEvent): + event.parallel_mode_run_id = parallel_mode_run_id + return event + if event.route_node_state.node_run_result: + metadata = event.route_node_state.node_run_result.metadata + if not metadata: + metadata = {} + if NodeRunMetadataKey.ITERATION_ID not in metadata: + metadata = { + **metadata, + NodeRunMetadataKey.ITERATION_ID: self.node_id, + NodeRunMetadataKey.PARALLEL_MODE_RUN_ID + if self.node_data.is_parallel + else NodeRunMetadataKey.ITERATION_INDEX: parallel_mode_run_id + if self.node_data.is_parallel + else iter_run_index, + } + event.route_node_state.node_run_result.metadata = metadata + return event + + def _run_single_iter( + self, + *, + iterator_list_value: Sequence[str], + variable_pool: VariablePool, + inputs: Mapping[str, list], + outputs: list, + start_at: datetime, + graph_engine: "GraphEngine", + iteration_graph: Graph, + iter_run_map: dict[str, float], + parallel_mode_run_id: Optional[str] = None, + ) -> Generator[NodeEvent | InNodeEvent, None, None]: + """ + run single iteration + """ + iter_start_at = datetime.now(UTC).replace(tzinfo=None) + + try: + rst = graph_engine.run() + # get current iteration index + index_variable = variable_pool.get([self.node_id, "index"]) + if not isinstance(index_variable, IntegerVariable): + raise IterationIndexNotFoundError(f"iteration {self.node_id} current index not found") + current_index = index_variable.value + iteration_run_id = parallel_mode_run_id if parallel_mode_run_id is not None else f"{current_index}" + next_index = int(current_index) + 1 + for event in rst: + if isinstance(event, (BaseNodeEvent | BaseParallelBranchEvent)) and not event.in_iteration_id: + event.in_iteration_id = self.node_id + + if ( + isinstance(event, BaseNodeEvent) + and event.node_type == NodeType.ITERATION_START + and not isinstance(event, NodeRunStreamChunkEvent) + ): + continue + + if isinstance(event, NodeRunSucceededEvent): + yield self._handle_event_metadata( + event=event, iter_run_index=current_index, parallel_mode_run_id=parallel_mode_run_id + ) + elif isinstance(event, BaseGraphEvent): + if isinstance(event, GraphRunFailedEvent): + # iteration run failed + if self.node_data.is_parallel: + yield IterationRunFailedEvent( + iteration_id=self.id, + iteration_node_id=self.node_id, + iteration_node_type=self.node_type, + iteration_node_data=self.node_data, + parallel_mode_run_id=parallel_mode_run_id, + start_at=start_at, + inputs=inputs, + outputs={"output": outputs}, + steps=len(iterator_list_value), + metadata={"total_tokens": graph_engine.graph_runtime_state.total_tokens}, + error=event.error, + ) + else: + yield IterationRunFailedEvent( + iteration_id=self.id, + iteration_node_id=self.node_id, + iteration_node_type=self.node_type, + iteration_node_data=self.node_data, + start_at=start_at, + inputs=inputs, + outputs={"output": outputs}, + steps=len(iterator_list_value), + metadata={"total_tokens": graph_engine.graph_runtime_state.total_tokens}, + error=event.error, + ) + yield RunCompletedEvent( + run_result=NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + error=event.error, + ) + ) + return + elif isinstance(event, InNodeEvent): + # event = cast(InNodeEvent, event) + metadata_event = self._handle_event_metadata( + event=event, iter_run_index=current_index, parallel_mode_run_id=parallel_mode_run_id + ) + if isinstance(event, NodeRunFailedEvent): + if self.node_data.error_handle_mode == ErrorHandleMode.CONTINUE_ON_ERROR: + yield NodeInIterationFailedEvent( + **metadata_event.model_dump(), + ) + outputs[current_index] = None + variable_pool.add([self.node_id, "index"], next_index) + if next_index < len(iterator_list_value): + variable_pool.add([self.node_id, "item"], iterator_list_value[next_index]) + duration = (datetime.now(UTC).replace(tzinfo=None) - iter_start_at).total_seconds() + iter_run_map[iteration_run_id] = duration + yield IterationRunNextEvent( + iteration_id=self.id, + iteration_node_id=self.node_id, + iteration_node_type=self.node_type, + iteration_node_data=self.node_data, + index=next_index, + parallel_mode_run_id=parallel_mode_run_id, + pre_iteration_output=None, + duration=duration, + ) + return + elif self.node_data.error_handle_mode == ErrorHandleMode.REMOVE_ABNORMAL_OUTPUT: + yield NodeInIterationFailedEvent( + **metadata_event.model_dump(), + ) + variable_pool.add([self.node_id, "index"], next_index) + + if next_index < len(iterator_list_value): + variable_pool.add([self.node_id, "item"], iterator_list_value[next_index]) + duration = (datetime.now(UTC).replace(tzinfo=None) - iter_start_at).total_seconds() + iter_run_map[iteration_run_id] = duration + yield IterationRunNextEvent( + iteration_id=self.id, + iteration_node_id=self.node_id, + iteration_node_type=self.node_type, + iteration_node_data=self.node_data, + index=next_index, + parallel_mode_run_id=parallel_mode_run_id, + pre_iteration_output=None, + duration=duration, + ) + return + elif self.node_data.error_handle_mode == ErrorHandleMode.TERMINATED: + yield IterationRunFailedEvent( + iteration_id=self.id, + iteration_node_id=self.node_id, + iteration_node_type=self.node_type, + iteration_node_data=self.node_data, + start_at=start_at, + inputs=inputs, + outputs={"output": None}, + steps=len(iterator_list_value), + metadata={"total_tokens": graph_engine.graph_runtime_state.total_tokens}, + error=event.error, + ) + yield metadata_event + + current_output_segment = variable_pool.get(self.node_data.output_selector) + if current_output_segment is None: + raise IterationNodeError("iteration output selector not found") + current_iteration_output = current_output_segment.value + outputs[current_index] = current_iteration_output + # remove all nodes outputs from variable pool + for node_id in iteration_graph.node_ids: + variable_pool.remove([node_id]) + + # move to next iteration + variable_pool.add([self.node_id, "index"], next_index) + + if next_index < len(iterator_list_value): + variable_pool.add([self.node_id, "item"], iterator_list_value[next_index]) + duration = (datetime.now(UTC).replace(tzinfo=None) - iter_start_at).total_seconds() + iter_run_map[iteration_run_id] = duration + yield IterationRunNextEvent( + iteration_id=self.id, + iteration_node_id=self.node_id, + iteration_node_type=self.node_type, + iteration_node_data=self.node_data, + index=next_index, + parallel_mode_run_id=parallel_mode_run_id, + pre_iteration_output=current_iteration_output or None, + duration=duration, + ) + + except IterationNodeError as e: + logger.warning(f"Iteration run failed:{str(e)}") + yield IterationRunFailedEvent( + iteration_id=self.id, + iteration_node_id=self.node_id, + iteration_node_type=self.node_type, + iteration_node_data=self.node_data, + start_at=start_at, + inputs=inputs, + outputs={"output": None}, + steps=len(iterator_list_value), + metadata={"total_tokens": graph_engine.graph_runtime_state.total_tokens}, + error=str(e), + ) + yield RunCompletedEvent( + run_result=NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + error=str(e), + ) + ) + + def _run_single_iter_parallel( + self, + *, + flask_app: Flask, + q: Queue, + iterator_list_value: Sequence[str], + inputs: Mapping[str, list], + outputs: list, + start_at: datetime, + graph_engine: "GraphEngine", + iteration_graph: Graph, + index: int, + item: Any, + iter_run_map: dict[str, float], + ): + """ + run single iteration in parallel mode + """ + with flask_app.app_context(): + parallel_mode_run_id = uuid.uuid4().hex + graph_engine_copy = graph_engine.create_copy() + variable_pool_copy = graph_engine_copy.graph_runtime_state.variable_pool + variable_pool_copy.add([self.node_id, "index"], index) + variable_pool_copy.add([self.node_id, "item"], item) + for event in self._run_single_iter( + iterator_list_value=iterator_list_value, + variable_pool=variable_pool_copy, + inputs=inputs, + outputs=outputs, + start_at=start_at, + graph_engine=graph_engine_copy, + iteration_graph=iteration_graph, + iter_run_map=iter_run_map, + parallel_mode_run_id=parallel_mode_run_id, + ): + q.put(event) + graph_engine.graph_runtime_state.total_tokens += graph_engine_copy.graph_runtime_state.total_tokens diff --git a/api/core/workflow/nodes/iteration/iteration_start_node.py b/api/core/workflow/nodes/iteration/iteration_start_node.py new file mode 100644 index 0000000000000000000000000000000000000000..6ab7c301066d9317cb966fba7f2d29246f783186 --- /dev/null +++ b/api/core/workflow/nodes/iteration/iteration_start_node.py @@ -0,0 +1,36 @@ +from collections.abc import Mapping, Sequence +from typing import Any + +from core.workflow.entities.node_entities import NodeRunResult +from core.workflow.nodes.base import BaseNode +from core.workflow.nodes.enums import NodeType +from core.workflow.nodes.iteration.entities import IterationNodeData, IterationStartNodeData +from models.workflow import WorkflowNodeExecutionStatus + + +class IterationStartNode(BaseNode): + """ + Iteration Start Node. + """ + + _node_data_cls = IterationStartNodeData + _node_type = NodeType.ITERATION_START + + def _run(self) -> NodeRunResult: + """ + Run the node. + """ + return NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED) + + @classmethod + def _extract_variable_selector_to_variable_mapping( + cls, graph_config: Mapping[str, Any], node_id: str, node_data: IterationNodeData + ) -> Mapping[str, Sequence[str]]: + """ + Extract variable selector to variable mapping + :param graph_config: graph config + :param node_id: node id + :param node_data: node data + :return: + """ + return {} diff --git a/api/core/workflow/nodes/knowledge_retrieval/__init__.py b/api/core/workflow/nodes/knowledge_retrieval/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4d4a4cbd9f13426ced1a51802325968549990781 --- /dev/null +++ b/api/core/workflow/nodes/knowledge_retrieval/__init__.py @@ -0,0 +1,3 @@ +from .knowledge_retrieval_node import KnowledgeRetrievalNode + +__all__ = ["KnowledgeRetrievalNode"] diff --git a/api/core/workflow/nodes/knowledge_retrieval/entities.py b/api/core/workflow/nodes/knowledge_retrieval/entities.py new file mode 100644 index 0000000000000000000000000000000000000000..e8972d1381d3ce3e405e50aaaeebfdd695faf516 --- /dev/null +++ b/api/core/workflow/nodes/knowledge_retrieval/entities.py @@ -0,0 +1,86 @@ +from typing import Any, Literal, Optional + +from pydantic import BaseModel + +from core.workflow.nodes.base import BaseNodeData + + +class RerankingModelConfig(BaseModel): + """ + Reranking Model Config. + """ + + provider: str + model: 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 MultipleRetrievalConfig(BaseModel): + """ + Multiple Retrieval Config. + """ + + top_k: int + score_threshold: Optional[float] = None + reranking_mode: str = "reranking_model" + reranking_enable: bool = True + reranking_model: Optional[RerankingModelConfig] = None + weights: Optional[WeightedScoreConfig] = None + + +class ModelConfig(BaseModel): + """ + Model Config. + """ + + provider: str + name: str + mode: str + completion_params: dict[str, Any] = {} + + +class SingleRetrievalConfig(BaseModel): + """ + Single Retrieval Config. + """ + + model: ModelConfig + + +class KnowledgeRetrievalNodeData(BaseNodeData): + """ + Knowledge retrieval Node Data. + """ + + type: str = "knowledge-retrieval" + query_variable_selector: list[str] + dataset_ids: list[str] + retrieval_mode: Literal["single", "multiple"] + multiple_retrieval_config: Optional[MultipleRetrievalConfig] = None + single_retrieval_config: Optional[SingleRetrievalConfig] = None diff --git a/api/core/workflow/nodes/knowledge_retrieval/exc.py b/api/core/workflow/nodes/knowledge_retrieval/exc.py new file mode 100644 index 0000000000000000000000000000000000000000..0c3b6e86fa37be0388455fe00160f3ddc90d382d --- /dev/null +++ b/api/core/workflow/nodes/knowledge_retrieval/exc.py @@ -0,0 +1,18 @@ +class KnowledgeRetrievalNodeError(ValueError): + """Base class for KnowledgeRetrievalNode errors.""" + + +class ModelNotExistError(KnowledgeRetrievalNodeError): + """Raised when the model does not exist.""" + + +class ModelCredentialsNotInitializedError(KnowledgeRetrievalNodeError): + """Raised when the model credentials are not initialized.""" + + +class ModelNotSupportedError(KnowledgeRetrievalNodeError): + """Raised when the model is not supported.""" + + +class ModelQuotaExceededError(KnowledgeRetrievalNodeError): + """Raised when the model provider quota is exceeded.""" diff --git a/api/core/workflow/nodes/knowledge_retrieval/knowledge_retrieval_node.py b/api/core/workflow/nodes/knowledge_retrieval/knowledge_retrieval_node.py new file mode 100644 index 0000000000000000000000000000000000000000..0f239af51ae79ccc8b98a056c261c6db3aace879 --- /dev/null +++ b/api/core/workflow/nodes/knowledge_retrieval/knowledge_retrieval_node.py @@ -0,0 +1,345 @@ +import logging +from collections.abc import Mapping, Sequence +from typing import Any, cast + +from sqlalchemy import func + +from core.app.app_config.entities import DatasetRetrieveConfigEntity +from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity +from core.entities.agent_entities import PlanningStrategy +from core.entities.model_entities import ModelStatus +from core.model_manager import ModelInstance, ModelManager +from core.model_runtime.entities.model_entities import ModelFeature, ModelType +from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel +from core.rag.datasource.retrieval_service import RetrievalService +from core.rag.retrieval.dataset_retrieval import DatasetRetrieval +from core.rag.retrieval.retrieval_methods import RetrievalMethod +from core.variables import StringSegment +from core.workflow.entities.node_entities import NodeRunResult +from core.workflow.nodes.base import BaseNode +from core.workflow.nodes.enums import NodeType +from extensions.ext_database import db +from models.dataset import Dataset, Document +from models.workflow import WorkflowNodeExecutionStatus + +from .entities import KnowledgeRetrievalNodeData +from .exc import ( + KnowledgeRetrievalNodeError, + ModelCredentialsNotInitializedError, + ModelNotExistError, + ModelNotSupportedError, + ModelQuotaExceededError, +) + +logger = logging.getLogger(__name__) + +default_retrieval_model = { + "search_method": RetrievalMethod.SEMANTIC_SEARCH.value, + "reranking_enable": False, + "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""}, + "top_k": 2, + "score_threshold_enabled": False, +} + + +class KnowledgeRetrievalNode(BaseNode[KnowledgeRetrievalNodeData]): + _node_data_cls = KnowledgeRetrievalNodeData + _node_type = NodeType.KNOWLEDGE_RETRIEVAL + + def _run(self) -> NodeRunResult: + # extract variables + variable = self.graph_runtime_state.variable_pool.get(self.node_data.query_variable_selector) + if not isinstance(variable, StringSegment): + return NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + inputs={}, + error="Query variable is not string type.", + ) + query = variable.value + variables = {"query": query} + if not query: + return NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, inputs=variables, error="Query is required." + ) + # retrieve knowledge + try: + results = self._fetch_dataset_retriever(node_data=self.node_data, query=query) + outputs = {"result": results} + return NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, inputs=variables, process_data=None, outputs=outputs + ) + + except KnowledgeRetrievalNodeError as e: + logger.warning("Error when running knowledge retrieval node") + return NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + inputs=variables, + error=str(e), + error_type=type(e).__name__, + ) + # Temporary handle all exceptions from DatasetRetrieval class here. + except Exception as e: + return NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + inputs=variables, + error=str(e), + error_type=type(e).__name__, + ) + + def _fetch_dataset_retriever(self, node_data: KnowledgeRetrievalNodeData, query: str) -> list[dict[str, Any]]: + available_datasets = [] + dataset_ids = node_data.dataset_ids + + # Subquery: Count the number of available documents for each dataset + subquery = ( + db.session.query(Document.dataset_id, func.count(Document.id).label("available_document_count")) + .filter( + Document.indexing_status == "completed", + Document.enabled == True, + Document.archived == False, + Document.dataset_id.in_(dataset_ids), + ) + .group_by(Document.dataset_id) + .having(func.count(Document.id) > 0) + .subquery() + ) + + results = ( + db.session.query(Dataset) + .outerjoin(subquery, Dataset.id == subquery.c.dataset_id) + .filter(Dataset.tenant_id == self.tenant_id, Dataset.id.in_(dataset_ids)) + .filter((subquery.c.available_document_count > 0) | (Dataset.provider == "external")) + .all() + ) + + for dataset in results: + # pass if dataset is not available + if not dataset: + continue + available_datasets.append(dataset) + all_documents = [] + dataset_retrieval = DatasetRetrieval() + if node_data.retrieval_mode == DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE.value: + # fetch model config + model_instance, model_config = self._fetch_model_config(node_data) + # check model is support tool calling + model_type_instance = model_config.provider_model_bundle.model_type_instance + model_type_instance = cast(LargeLanguageModel, model_type_instance) + # get model schema + model_schema = model_type_instance.get_model_schema( + model=model_config.model, credentials=model_config.credentials + ) + + if model_schema: + planning_strategy = PlanningStrategy.REACT_ROUTER + features = model_schema.features + if features: + if ModelFeature.TOOL_CALL in features or ModelFeature.MULTI_TOOL_CALL in features: + planning_strategy = PlanningStrategy.ROUTER + all_documents = dataset_retrieval.single_retrieve( + available_datasets=available_datasets, + tenant_id=self.tenant_id, + user_id=self.user_id, + app_id=self.app_id, + user_from=self.user_from.value, + query=query, + model_config=model_config, + model_instance=model_instance, + planning_strategy=planning_strategy, + ) + elif node_data.retrieval_mode == DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE.value: + if node_data.multiple_retrieval_config is None: + raise ValueError("multiple_retrieval_config is required") + if node_data.multiple_retrieval_config.reranking_mode == "reranking_model": + if node_data.multiple_retrieval_config.reranking_model: + reranking_model = { + "reranking_provider_name": node_data.multiple_retrieval_config.reranking_model.provider, + "reranking_model_name": node_data.multiple_retrieval_config.reranking_model.model, + } + else: + reranking_model = None + weights = None + elif node_data.multiple_retrieval_config.reranking_mode == "weighted_score": + if node_data.multiple_retrieval_config.weights is None: + raise ValueError("weights is required") + reranking_model = None + vector_setting = node_data.multiple_retrieval_config.weights.vector_setting + weights = { + "vector_setting": { + "vector_weight": vector_setting.vector_weight, + "embedding_provider_name": vector_setting.embedding_provider_name, + "embedding_model_name": vector_setting.embedding_model_name, + }, + "keyword_setting": { + "keyword_weight": node_data.multiple_retrieval_config.weights.keyword_setting.keyword_weight + }, + } + else: + reranking_model = None + weights = None + all_documents = dataset_retrieval.multiple_retrieve( + app_id=self.app_id, + tenant_id=self.tenant_id, + user_id=self.user_id, + user_from=self.user_from.value, + available_datasets=available_datasets, + query=query, + top_k=node_data.multiple_retrieval_config.top_k, + score_threshold=node_data.multiple_retrieval_config.score_threshold + if node_data.multiple_retrieval_config.score_threshold is not None + else 0.0, + reranking_mode=node_data.multiple_retrieval_config.reranking_mode, + reranking_model=reranking_model, + weights=weights, + reranking_enable=node_data.multiple_retrieval_config.reranking_enable, + ) + dify_documents = [item for item in all_documents if item.provider == "dify"] + external_documents = [item for item in all_documents if item.provider == "external"] + retrieval_resource_list = [] + # deal with external documents + for item in external_documents: + source = { + "metadata": { + "_source": "knowledge", + "dataset_id": item.metadata.get("dataset_id"), + "dataset_name": item.metadata.get("dataset_name"), + "document_name": item.metadata.get("title"), + "data_source_type": "external", + "retriever_from": "workflow", + "score": item.metadata.get("score"), + }, + "title": item.metadata.get("title"), + "content": item.page_content, + } + retrieval_resource_list.append(source) + # deal with dify documents + if dify_documents: + records = RetrievalService.format_retrieval_documents(dify_documents) + if records: + for record in records: + segment = record.segment + dataset = Dataset.query.filter_by(id=segment.dataset_id).first() + document = Document.query.filter( + Document.id == segment.document_id, + Document.enabled == True, + Document.archived == False, + ).first() + if dataset and document: + source = { + "metadata": { + "_source": "knowledge", + "dataset_id": dataset.id, + "dataset_name": dataset.name, + "document_id": document.id, + "document_name": document.name, + "document_data_source_type": document.data_source_type, + "segment_id": segment.id, + "retriever_from": "workflow", + "score": record.score or 0.0, + "segment_hit_count": segment.hit_count, + "segment_word_count": segment.word_count, + "segment_position": segment.position, + "segment_index_node_hash": segment.index_node_hash, + }, + "title": document.name, + } + if segment.answer: + source["content"] = f"question:{segment.get_sign_content()} \nanswer:{segment.answer}" + else: + source["content"] = segment.get_sign_content() + retrieval_resource_list.append(source) + if retrieval_resource_list: + retrieval_resource_list = sorted( + retrieval_resource_list, + key=lambda x: x["metadata"]["score"] if x["metadata"].get("score") is not None else 0.0, + reverse=True, + ) + for position, item in enumerate(retrieval_resource_list, start=1): + item["metadata"]["position"] = position + return retrieval_resource_list + + @classmethod + def _extract_variable_selector_to_variable_mapping( + cls, + *, + graph_config: Mapping[str, Any], + node_id: str, + node_data: KnowledgeRetrievalNodeData, + ) -> Mapping[str, Sequence[str]]: + """ + Extract variable selector to variable mapping + :param graph_config: graph config + :param node_id: node id + :param node_data: node data + :return: + """ + variable_mapping = {} + variable_mapping[node_id + ".query"] = node_data.query_variable_selector + return variable_mapping + + def _fetch_model_config( + self, node_data: KnowledgeRetrievalNodeData + ) -> tuple[ModelInstance, ModelConfigWithCredentialsEntity]: + """ + Fetch model config + :param node_data: node data + :return: + """ + if node_data.single_retrieval_config is None: + raise ValueError("single_retrieval_config is required") + model_name = node_data.single_retrieval_config.model.name + provider_name = node_data.single_retrieval_config.model.provider + + model_manager = ModelManager() + model_instance = model_manager.get_model_instance( + tenant_id=self.tenant_id, model_type=ModelType.LLM, provider=provider_name, model=model_name + ) + + provider_model_bundle = model_instance.provider_model_bundle + model_type_instance = model_instance.model_type_instance + model_type_instance = cast(LargeLanguageModel, model_type_instance) + + model_credentials = model_instance.credentials + + # check model + provider_model = provider_model_bundle.configuration.get_provider_model( + model=model_name, model_type=ModelType.LLM + ) + + if provider_model is None: + raise ModelNotExistError(f"Model {model_name} not exist.") + + if provider_model.status == ModelStatus.NO_CONFIGURE: + raise ModelCredentialsNotInitializedError(f"Model {model_name} credentials is not initialized.") + elif provider_model.status == ModelStatus.NO_PERMISSION: + raise ModelNotSupportedError(f"Dify Hosted OpenAI {model_name} currently not support.") + elif provider_model.status == ModelStatus.QUOTA_EXCEEDED: + raise ModelQuotaExceededError(f"Model provider {provider_name} quota exceeded.") + + # model config + completion_params = node_data.single_retrieval_config.model.completion_params + stop = [] + if "stop" in completion_params: + stop = completion_params["stop"] + del completion_params["stop"] + + # get model mode + model_mode = node_data.single_retrieval_config.model.mode + if not model_mode: + raise ModelNotExistError("LLM mode is required.") + + model_schema = model_type_instance.get_model_schema(model_name, model_credentials) + + if not model_schema: + raise ModelNotExistError(f"Model {model_name} not exist.") + + return model_instance, ModelConfigWithCredentialsEntity( + provider=provider_name, + model=model_name, + model_schema=model_schema, + mode=model_mode, + provider_model_bundle=provider_model_bundle, + credentials=model_credentials, + parameters=completion_params, + stop=stop, + ) diff --git a/api/core/workflow/nodes/list_operator/__init__.py b/api/core/workflow/nodes/list_operator/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1877586ef41145fec681290932d986f393f1ea25 --- /dev/null +++ b/api/core/workflow/nodes/list_operator/__init__.py @@ -0,0 +1,3 @@ +from .node import ListOperatorNode + +__all__ = ["ListOperatorNode"] diff --git a/api/core/workflow/nodes/list_operator/entities.py b/api/core/workflow/nodes/list_operator/entities.py new file mode 100644 index 0000000000000000000000000000000000000000..75df784a9222269cdafd670878cb1bc0afcca65f --- /dev/null +++ b/api/core/workflow/nodes/list_operator/entities.py @@ -0,0 +1,62 @@ +from collections.abc import Sequence +from typing import Literal + +from pydantic import BaseModel, Field + +from core.workflow.nodes.base import BaseNodeData + +_Condition = Literal[ + # string conditions + "contains", + "start with", + "end with", + "is", + "in", + "empty", + "not contains", + "is not", + "not in", + "not empty", + # number conditions + "=", + "≠", + "<", + ">", + "≥", + "≤", +] + + +class FilterCondition(BaseModel): + key: str = "" + comparison_operator: _Condition = "contains" + value: str | Sequence[str] = "" + + +class FilterBy(BaseModel): + enabled: bool = False + conditions: Sequence[FilterCondition] = Field(default_factory=list) + + +class OrderBy(BaseModel): + enabled: bool = False + key: str = "" + value: Literal["asc", "desc"] = "asc" + + +class Limit(BaseModel): + enabled: bool = False + size: int = -1 + + +class ExtractConfig(BaseModel): + enabled: bool = False + serial: str = "1" + + +class ListOperatorNodeData(BaseNodeData): + variable: Sequence[str] = Field(default_factory=list) + filter_by: FilterBy + order_by: OrderBy + limit: Limit + extract_by: ExtractConfig = Field(default_factory=ExtractConfig) diff --git a/api/core/workflow/nodes/list_operator/exc.py b/api/core/workflow/nodes/list_operator/exc.py new file mode 100644 index 0000000000000000000000000000000000000000..f88aa0be29c92aeb7ad89bf3a764bf072be94533 --- /dev/null +++ b/api/core/workflow/nodes/list_operator/exc.py @@ -0,0 +1,16 @@ +class ListOperatorError(ValueError): + """Base class for all ListOperator errors.""" + + pass + + +class InvalidFilterValueError(ListOperatorError): + pass + + +class InvalidKeyError(ListOperatorError): + pass + + +class InvalidConditionError(ListOperatorError): + pass diff --git a/api/core/workflow/nodes/list_operator/node.py b/api/core/workflow/nodes/list_operator/node.py new file mode 100644 index 0000000000000000000000000000000000000000..432c57294ecbe97ecad057f205d66bd1e060aec0 --- /dev/null +++ b/api/core/workflow/nodes/list_operator/node.py @@ -0,0 +1,316 @@ +from collections.abc import Callable, Sequence +from typing import Any, Literal, Union + +from core.file import File +from core.variables import ArrayFileSegment, ArrayNumberSegment, ArrayStringSegment +from core.workflow.entities.node_entities import NodeRunResult +from core.workflow.nodes.base import BaseNode +from core.workflow.nodes.enums import NodeType +from models.workflow import WorkflowNodeExecutionStatus + +from .entities import ListOperatorNodeData +from .exc import InvalidConditionError, InvalidFilterValueError, InvalidKeyError, ListOperatorError + + +class ListOperatorNode(BaseNode[ListOperatorNodeData]): + _node_data_cls = ListOperatorNodeData + _node_type = NodeType.LIST_OPERATOR + + def _run(self): + inputs: dict[str, list] = {} + process_data: dict[str, list] = {} + outputs: dict[str, Any] = {} + + variable = self.graph_runtime_state.variable_pool.get(self.node_data.variable) + if variable is None: + error_message = f"Variable not found for selector: {self.node_data.variable}" + return NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, error=error_message, inputs=inputs, outputs=outputs + ) + if not variable.value: + inputs = {"variable": []} + process_data = {"variable": []} + outputs = {"result": [], "first_record": None, "last_record": None} + return NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + inputs=inputs, + process_data=process_data, + outputs=outputs, + ) + if not isinstance(variable, ArrayFileSegment | ArrayNumberSegment | ArrayStringSegment): + error_message = ( + f"Variable {self.node_data.variable} is not an ArrayFileSegment, ArrayNumberSegment " + "or ArrayStringSegment" + ) + return NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, error=error_message, inputs=inputs, outputs=outputs + ) + + if isinstance(variable, ArrayFileSegment): + inputs = {"variable": [item.to_dict() for item in variable.value]} + process_data["variable"] = [item.to_dict() for item in variable.value] + else: + inputs = {"variable": variable.value} + process_data["variable"] = variable.value + + try: + # Filter + if self.node_data.filter_by.enabled: + variable = self._apply_filter(variable) + + # Extract + if self.node_data.extract_by.enabled: + variable = self._extract_slice(variable) + + # Order + if self.node_data.order_by.enabled: + variable = self._apply_order(variable) + + # Slice + if self.node_data.limit.enabled: + variable = self._apply_slice(variable) + + outputs = { + "result": variable.value, + "first_record": variable.value[0] if variable.value else None, + "last_record": variable.value[-1] if variable.value else None, + } + return NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + inputs=inputs, + process_data=process_data, + outputs=outputs, + ) + except ListOperatorError as e: + return NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + error=str(e), + inputs=inputs, + process_data=process_data, + outputs=outputs, + ) + + def _apply_filter( + self, variable: Union[ArrayFileSegment, ArrayNumberSegment, ArrayStringSegment] + ) -> Union[ArrayFileSegment, ArrayNumberSegment, ArrayStringSegment]: + filter_func: Callable[[Any], bool] + result: list[Any] = [] + for condition in self.node_data.filter_by.conditions: + if isinstance(variable, ArrayStringSegment): + if not isinstance(condition.value, str): + raise InvalidFilterValueError(f"Invalid filter value: {condition.value}") + value = self.graph_runtime_state.variable_pool.convert_template(condition.value).text + filter_func = _get_string_filter_func(condition=condition.comparison_operator, value=value) + result = list(filter(filter_func, variable.value)) + variable = variable.model_copy(update={"value": result}) + elif isinstance(variable, ArrayNumberSegment): + if not isinstance(condition.value, str): + raise InvalidFilterValueError(f"Invalid filter value: {condition.value}") + value = self.graph_runtime_state.variable_pool.convert_template(condition.value).text + filter_func = _get_number_filter_func(condition=condition.comparison_operator, value=float(value)) + result = list(filter(filter_func, variable.value)) + variable = variable.model_copy(update={"value": result}) + elif isinstance(variable, ArrayFileSegment): + if isinstance(condition.value, str): + value = self.graph_runtime_state.variable_pool.convert_template(condition.value).text + else: + value = condition.value + filter_func = _get_file_filter_func( + key=condition.key, + condition=condition.comparison_operator, + value=value, + ) + result = list(filter(filter_func, variable.value)) + variable = variable.model_copy(update={"value": result}) + return variable + + def _apply_order( + self, variable: Union[ArrayFileSegment, ArrayNumberSegment, ArrayStringSegment] + ) -> Union[ArrayFileSegment, ArrayNumberSegment, ArrayStringSegment]: + if isinstance(variable, ArrayStringSegment): + result = _order_string(order=self.node_data.order_by.value, array=variable.value) + variable = variable.model_copy(update={"value": result}) + elif isinstance(variable, ArrayNumberSegment): + result = _order_number(order=self.node_data.order_by.value, array=variable.value) + variable = variable.model_copy(update={"value": result}) + elif isinstance(variable, ArrayFileSegment): + result = _order_file( + order=self.node_data.order_by.value, order_by=self.node_data.order_by.key, array=variable.value + ) + variable = variable.model_copy(update={"value": result}) + return variable + + def _apply_slice( + self, variable: Union[ArrayFileSegment, ArrayNumberSegment, ArrayStringSegment] + ) -> Union[ArrayFileSegment, ArrayNumberSegment, ArrayStringSegment]: + result = variable.value[: self.node_data.limit.size] + return variable.model_copy(update={"value": result}) + + def _extract_slice( + self, variable: Union[ArrayFileSegment, ArrayNumberSegment, ArrayStringSegment] + ) -> Union[ArrayFileSegment, ArrayNumberSegment, ArrayStringSegment]: + value = int(self.graph_runtime_state.variable_pool.convert_template(self.node_data.extract_by.serial).text) - 1 + if len(variable.value) > int(value): + result = variable.value[value] + else: + result = "" + return variable.model_copy(update={"value": [result]}) + + +def _get_file_extract_number_func(*, key: str) -> Callable[[File], int]: + match key: + case "size": + return lambda x: x.size + case _: + raise InvalidKeyError(f"Invalid key: {key}") + + +def _get_file_extract_string_func(*, key: str) -> Callable[[File], str]: + match key: + case "name": + return lambda x: x.filename or "" + case "type": + return lambda x: x.type + case "extension": + return lambda x: x.extension or "" + case "mime_type": + return lambda x: x.mime_type or "" + case "transfer_method": + return lambda x: x.transfer_method + case "url": + return lambda x: x.remote_url or "" + case _: + raise InvalidKeyError(f"Invalid key: {key}") + + +def _get_string_filter_func(*, condition: str, value: str) -> Callable[[str], bool]: + match condition: + case "contains": + return _contains(value) + case "start with": + return _startswith(value) + case "end with": + return _endswith(value) + case "is": + return _is(value) + case "in": + return _in(value) + case "empty": + return lambda x: x == "" + case "not contains": + return lambda x: not _contains(value)(x) + case "is not": + return lambda x: not _is(value)(x) + case "not in": + return lambda x: not _in(value)(x) + case "not empty": + return lambda x: x != "" + case _: + raise InvalidConditionError(f"Invalid condition: {condition}") + + +def _get_sequence_filter_func(*, condition: str, value: Sequence[str]) -> Callable[[str], bool]: + match condition: + case "in": + return _in(value) + case "not in": + return lambda x: not _in(value)(x) + case _: + raise InvalidConditionError(f"Invalid condition: {condition}") + + +def _get_number_filter_func(*, condition: str, value: int | float) -> Callable[[int | float], bool]: + match condition: + case "=": + return _eq(value) + case "≠": + return _ne(value) + case "<": + return _lt(value) + case "≤": + return _le(value) + case ">": + return _gt(value) + case "≥": + return _ge(value) + case _: + raise InvalidConditionError(f"Invalid condition: {condition}") + + +def _get_file_filter_func(*, key: str, condition: str, value: str | Sequence[str]) -> Callable[[File], bool]: + extract_func: Callable[[File], Any] + if key in {"name", "extension", "mime_type", "url"} and isinstance(value, str): + extract_func = _get_file_extract_string_func(key=key) + return lambda x: _get_string_filter_func(condition=condition, value=value)(extract_func(x)) + if key in {"type", "transfer_method"} and isinstance(value, Sequence): + extract_func = _get_file_extract_string_func(key=key) + return lambda x: _get_sequence_filter_func(condition=condition, value=value)(extract_func(x)) + elif key == "size" and isinstance(value, str): + extract_func = _get_file_extract_number_func(key=key) + return lambda x: _get_number_filter_func(condition=condition, value=float(value))(extract_func(x)) + else: + raise InvalidKeyError(f"Invalid key: {key}") + + +def _contains(value: str) -> Callable[[str], bool]: + return lambda x: value in x + + +def _startswith(value: str) -> Callable[[str], bool]: + return lambda x: x.startswith(value) + + +def _endswith(value: str) -> Callable[[str], bool]: + return lambda x: x.endswith(value) + + +def _is(value: str) -> Callable[[str], bool]: + return lambda x: x is value + + +def _in(value: str | Sequence[str]) -> Callable[[str], bool]: + return lambda x: x in value + + +def _eq(value: int | float) -> Callable[[int | float], bool]: + return lambda x: x == value + + +def _ne(value: int | float) -> Callable[[int | float], bool]: + return lambda x: x != value + + +def _lt(value: int | float) -> Callable[[int | float], bool]: + return lambda x: x < value + + +def _le(value: int | float) -> Callable[[int | float], bool]: + return lambda x: x <= value + + +def _gt(value: int | float) -> Callable[[int | float], bool]: + return lambda x: x > value + + +def _ge(value: int | float) -> Callable[[int | float], bool]: + return lambda x: x >= value + + +def _order_number(*, order: Literal["asc", "desc"], array: Sequence[int | float]): + return sorted(array, key=lambda x: x, reverse=order == "desc") + + +def _order_string(*, order: Literal["asc", "desc"], array: Sequence[str]): + return sorted(array, key=lambda x: x, reverse=order == "desc") + + +def _order_file(*, order: Literal["asc", "desc"], order_by: str = "", array: Sequence[File]): + extract_func: Callable[[File], Any] + if order_by in {"name", "type", "extension", "mime_type", "transfer_method", "url"}: + extract_func = _get_file_extract_string_func(key=order_by) + return sorted(array, key=lambda x: extract_func(x), reverse=order == "desc") + elif order_by == "size": + extract_func = _get_file_extract_number_func(key=order_by) + return sorted(array, key=lambda x: extract_func(x), reverse=order == "desc") + else: + raise InvalidKeyError(f"Invalid order key: {order_by}") diff --git a/api/core/workflow/nodes/llm/__init__.py b/api/core/workflow/nodes/llm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f7bc713f63174efb5f891f7af54fda3fc5ee9931 --- /dev/null +++ b/api/core/workflow/nodes/llm/__init__.py @@ -0,0 +1,17 @@ +from .entities import ( + LLMNodeChatModelMessage, + LLMNodeCompletionModelPromptTemplate, + LLMNodeData, + ModelConfig, + VisionConfig, +) +from .node import LLMNode + +__all__ = [ + "LLMNode", + "LLMNodeChatModelMessage", + "LLMNodeCompletionModelPromptTemplate", + "LLMNodeData", + "ModelConfig", + "VisionConfig", +] diff --git a/api/core/workflow/nodes/llm/entities.py b/api/core/workflow/nodes/llm/entities.py new file mode 100644 index 0000000000000000000000000000000000000000..bf54fdb80c630f86332f791507da36a01bb03846 --- /dev/null +++ b/api/core/workflow/nodes/llm/entities.py @@ -0,0 +1,74 @@ +from collections.abc import Sequence +from typing import Any, Optional + +from pydantic import BaseModel, Field, field_validator + +from core.model_runtime.entities import ImagePromptMessageContent, LLMMode +from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate, MemoryConfig +from core.workflow.entities.variable_entities import VariableSelector +from core.workflow.nodes.base import BaseNodeData + + +class ModelConfig(BaseModel): + provider: str + name: str + mode: LLMMode + completion_params: dict[str, Any] = {} + + +class ContextConfig(BaseModel): + enabled: bool + variable_selector: Optional[list[str]] = None + + +class VisionConfigOptions(BaseModel): + variable_selector: Sequence[str] = Field(default_factory=lambda: ["sys", "files"]) + detail: ImagePromptMessageContent.DETAIL = ImagePromptMessageContent.DETAIL.HIGH + + +class VisionConfig(BaseModel): + enabled: bool = False + configs: VisionConfigOptions = Field(default_factory=VisionConfigOptions) + + @field_validator("configs", mode="before") + @classmethod + def convert_none_configs(cls, v: Any): + if v is None: + return VisionConfigOptions() + return v + + +class PromptConfig(BaseModel): + jinja2_variables: Sequence[VariableSelector] = Field(default_factory=list) + + @field_validator("jinja2_variables", mode="before") + @classmethod + def convert_none_jinja2_variables(cls, v: Any): + if v is None: + return [] + return v + + +class LLMNodeChatModelMessage(ChatModelMessage): + text: str = "" + jinja2_text: Optional[str] = None + + +class LLMNodeCompletionModelPromptTemplate(CompletionModelPromptTemplate): + jinja2_text: Optional[str] = None + + +class LLMNodeData(BaseNodeData): + model: ModelConfig + prompt_template: Sequence[LLMNodeChatModelMessage] | LLMNodeCompletionModelPromptTemplate + prompt_config: PromptConfig = Field(default_factory=PromptConfig) + memory: Optional[MemoryConfig] = None + context: ContextConfig + vision: VisionConfig = Field(default_factory=VisionConfig) + + @field_validator("prompt_config", mode="before") + @classmethod + def convert_none_prompt_config(cls, v: Any): + if v is None: + return PromptConfig() + return v diff --git a/api/core/workflow/nodes/llm/exc.py b/api/core/workflow/nodes/llm/exc.py new file mode 100644 index 0000000000000000000000000000000000000000..6599221691bfac2666053e5026828827365b1288 --- /dev/null +++ b/api/core/workflow/nodes/llm/exc.py @@ -0,0 +1,40 @@ +class LLMNodeError(ValueError): + """Base class for LLM Node errors.""" + + +class VariableNotFoundError(LLMNodeError): + """Raised when a required variable is not found.""" + + +class InvalidContextStructureError(LLMNodeError): + """Raised when the context structure is invalid.""" + + +class InvalidVariableTypeError(LLMNodeError): + """Raised when the variable type is invalid.""" + + +class ModelNotExistError(LLMNodeError): + """Raised when the specified model does not exist.""" + + +class LLMModeRequiredError(LLMNodeError): + """Raised when LLM mode is required but not provided.""" + + +class NoPromptFoundError(LLMNodeError): + """Raised when no prompt is found in the LLM configuration.""" + + +class TemplateTypeNotSupportError(LLMNodeError): + def __init__(self, *, type_name: str): + super().__init__(f"Prompt type {type_name} is not supported.") + + +class MemoryRolePrefixRequiredError(LLMNodeError): + """Raised when memory role prefix is required for completion model.""" + + +class FileTypeNotSupportError(LLMNodeError): + def __init__(self, *, type_name: str): + super().__init__(f"{type_name} type is not supported by this model") diff --git a/api/core/workflow/nodes/llm/node.py b/api/core/workflow/nodes/llm/node.py new file mode 100644 index 0000000000000000000000000000000000000000..7e28aa7a3ffb3d536a4948a3db1b45fd2f67d9b5 --- /dev/null +++ b/api/core/workflow/nodes/llm/node.py @@ -0,0 +1,1040 @@ +import json +import logging +from collections.abc import Generator, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Optional, cast + +from configs import dify_config +from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity +from core.entities.model_entities import ModelStatus +from core.entities.provider_entities import QuotaUnit +from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError +from core.file import FileType, file_manager +from core.helper.code_executor import CodeExecutor, CodeLanguage +from core.memory.token_buffer_memory import TokenBufferMemory +from core.model_manager import ModelInstance, ModelManager +from core.model_runtime.entities import ( + ImagePromptMessageContent, + PromptMessage, + PromptMessageContentType, + TextPromptMessageContent, +) +from core.model_runtime.entities.llm_entities import LLMResult, LLMUsage +from core.model_runtime.entities.message_entities import ( + AssistantPromptMessage, + PromptMessageContent, + PromptMessageRole, + SystemPromptMessage, + UserPromptMessage, +) +from core.model_runtime.entities.model_entities import ModelFeature, ModelPropertyKey, ModelType +from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel +from core.model_runtime.utils.encoders import jsonable_encoder +from core.prompt.entities.advanced_prompt_entities import CompletionModelPromptTemplate, MemoryConfig +from core.prompt.utils.prompt_message_util import PromptMessageUtil +from core.variables import ( + ArrayAnySegment, + ArrayFileSegment, + ArraySegment, + FileSegment, + NoneSegment, + ObjectSegment, + StringSegment, +) +from core.workflow.constants import SYSTEM_VARIABLE_NODE_ID +from core.workflow.entities.node_entities import NodeRunMetadataKey, NodeRunResult +from core.workflow.entities.variable_entities import VariableSelector +from core.workflow.entities.variable_pool import VariablePool +from core.workflow.enums import SystemVariableKey +from core.workflow.graph_engine.entities.event import InNodeEvent +from core.workflow.nodes.base import BaseNode +from core.workflow.nodes.enums import NodeType +from core.workflow.nodes.event import ( + ModelInvokeCompletedEvent, + NodeEvent, + RunCompletedEvent, + RunRetrieverResourceEvent, + RunStreamChunkEvent, +) +from core.workflow.utils.variable_template_parser import VariableTemplateParser +from extensions.ext_database import db +from models.model import Conversation +from models.provider import Provider, ProviderType +from models.workflow import WorkflowNodeExecutionStatus + +from .entities import ( + LLMNodeChatModelMessage, + LLMNodeCompletionModelPromptTemplate, + LLMNodeData, + ModelConfig, +) +from .exc import ( + InvalidContextStructureError, + InvalidVariableTypeError, + LLMModeRequiredError, + LLMNodeError, + MemoryRolePrefixRequiredError, + ModelNotExistError, + NoPromptFoundError, + TemplateTypeNotSupportError, + VariableNotFoundError, +) + +if TYPE_CHECKING: + from core.file.models import File + +logger = logging.getLogger(__name__) + + +class LLMNode(BaseNode[LLMNodeData]): + _node_data_cls = LLMNodeData + _node_type = NodeType.LLM + + def _run(self) -> Generator[NodeEvent | InNodeEvent, None, None]: + node_inputs: Optional[dict[str, Any]] = None + process_data = None + + try: + # init messages template + self.node_data.prompt_template = self._transform_chat_messages(self.node_data.prompt_template) + + # fetch variables and fetch values from variable pool + inputs = self._fetch_inputs(node_data=self.node_data) + + # fetch jinja2 inputs + jinja_inputs = self._fetch_jinja_inputs(node_data=self.node_data) + + # merge inputs + inputs.update(jinja_inputs) + + node_inputs = {} + + # fetch files + files = ( + self._fetch_files(selector=self.node_data.vision.configs.variable_selector) + if self.node_data.vision.enabled + else [] + ) + + if files: + node_inputs["#files#"] = [file.to_dict() for file in files] + + # fetch context value + generator = self._fetch_context(node_data=self.node_data) + context = None + for event in generator: + if isinstance(event, RunRetrieverResourceEvent): + context = event.context + yield event + + if context: + node_inputs["#context#"] = context + + # fetch model config + model_instance, model_config = self._fetch_model_config(self.node_data.model) + + # fetch memory + memory = self._fetch_memory(node_data_memory=self.node_data.memory, model_instance=model_instance) + + query = None + if self.node_data.memory: + query = self.node_data.memory.query_prompt_template + if not query and ( + query_variable := self.graph_runtime_state.variable_pool.get( + (SYSTEM_VARIABLE_NODE_ID, SystemVariableKey.QUERY) + ) + ): + query = query_variable.text + + prompt_messages, stop = self._fetch_prompt_messages( + sys_query=query, + sys_files=files, + context=context, + memory=memory, + model_config=model_config, + prompt_template=self.node_data.prompt_template, + memory_config=self.node_data.memory, + vision_enabled=self.node_data.vision.enabled, + vision_detail=self.node_data.vision.configs.detail, + variable_pool=self.graph_runtime_state.variable_pool, + jinja2_variables=self.node_data.prompt_config.jinja2_variables, + ) + + process_data = { + "model_mode": model_config.mode, + "prompts": PromptMessageUtil.prompt_messages_to_prompt_for_saving( + model_mode=model_config.mode, prompt_messages=prompt_messages + ), + "model_provider": model_config.provider, + "model_name": model_config.model, + } + + # handle invoke result + generator = self._invoke_llm( + node_data_model=self.node_data.model, + model_instance=model_instance, + prompt_messages=prompt_messages, + stop=stop, + ) + + result_text = "" + usage = LLMUsage.empty_usage() + finish_reason = None + for event in generator: + if isinstance(event, RunStreamChunkEvent): + yield event + elif isinstance(event, ModelInvokeCompletedEvent): + result_text = event.text + usage = event.usage + finish_reason = event.finish_reason + # deduct quota + self.deduct_llm_quota(tenant_id=self.tenant_id, model_instance=model_instance, usage=usage) + break + except LLMNodeError as e: + yield RunCompletedEvent( + run_result=NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + error=str(e), + inputs=node_inputs, + process_data=process_data, + error_type=type(e).__name__, + ) + ) + except Exception as e: + yield RunCompletedEvent( + run_result=NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + error=str(e), + inputs=node_inputs, + process_data=process_data, + ) + ) + + outputs = {"text": result_text, "usage": jsonable_encoder(usage), "finish_reason": finish_reason} + + yield RunCompletedEvent( + run_result=NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + inputs=node_inputs, + process_data=process_data, + outputs=outputs, + metadata={ + NodeRunMetadataKey.TOTAL_TOKENS: usage.total_tokens, + NodeRunMetadataKey.TOTAL_PRICE: usage.total_price, + NodeRunMetadataKey.CURRENCY: usage.currency, + }, + llm_usage=usage, + ) + ) + + def _invoke_llm( + self, + node_data_model: ModelConfig, + model_instance: ModelInstance, + prompt_messages: Sequence[PromptMessage], + stop: Optional[Sequence[str]] = None, + ) -> Generator[NodeEvent, None, None]: + db.session.close() + + invoke_result = model_instance.invoke_llm( + prompt_messages=prompt_messages, + model_parameters=node_data_model.completion_params, + stop=stop, + stream=True, + user=self.user_id, + ) + + return self._handle_invoke_result(invoke_result=invoke_result) + + def _handle_invoke_result(self, invoke_result: LLMResult | Generator) -> Generator[NodeEvent, None, None]: + if isinstance(invoke_result, LLMResult): + return + + model = None + prompt_messages: list[PromptMessage] = [] + full_text = "" + usage = None + finish_reason = None + for result in invoke_result: + text = result.delta.message.content + full_text += text + + yield RunStreamChunkEvent(chunk_content=text, from_variable_selector=[self.node_id, "text"]) + + if not model: + model = result.model + + if not prompt_messages: + prompt_messages = result.prompt_messages + + if not usage and result.delta.usage: + usage = result.delta.usage + + if not finish_reason and result.delta.finish_reason: + finish_reason = result.delta.finish_reason + + if not usage: + usage = LLMUsage.empty_usage() + + yield ModelInvokeCompletedEvent(text=full_text, usage=usage, finish_reason=finish_reason) + + def _transform_chat_messages( + self, messages: Sequence[LLMNodeChatModelMessage] | LLMNodeCompletionModelPromptTemplate, / + ) -> Sequence[LLMNodeChatModelMessage] | LLMNodeCompletionModelPromptTemplate: + if isinstance(messages, LLMNodeCompletionModelPromptTemplate): + if messages.edition_type == "jinja2" and messages.jinja2_text: + messages.text = messages.jinja2_text + + return messages + + for message in messages: + if message.edition_type == "jinja2" and message.jinja2_text: + message.text = message.jinja2_text + + return messages + + def _fetch_jinja_inputs(self, node_data: LLMNodeData) -> dict[str, str]: + variables: dict[str, Any] = {} + + if not node_data.prompt_config: + return variables + + for variable_selector in node_data.prompt_config.jinja2_variables or []: + variable_name = variable_selector.variable + variable = self.graph_runtime_state.variable_pool.get(variable_selector.value_selector) + if variable is None: + raise VariableNotFoundError(f"Variable {variable_selector.variable} not found") + + def parse_dict(input_dict: Mapping[str, Any]) -> str: + """ + Parse dict into string + """ + # check if it's a context structure + if "metadata" in input_dict and "_source" in input_dict["metadata"] and "content" in input_dict: + return str(input_dict["content"]) + + # else, parse the dict + try: + return json.dumps(input_dict, ensure_ascii=False) + except Exception: + return str(input_dict) + + if isinstance(variable, ArraySegment): + result = "" + for item in variable.value: + if isinstance(item, dict): + result += parse_dict(item) + else: + result += str(item) + result += "\n" + value = result.strip() + elif isinstance(variable, ObjectSegment): + value = parse_dict(variable.value) + else: + value = variable.text + + variables[variable_name] = value + + return variables + + def _fetch_inputs(self, node_data: LLMNodeData) -> dict[str, Any]: + inputs = {} + prompt_template = node_data.prompt_template + + variable_selectors = [] + if isinstance(prompt_template, list): + for prompt in prompt_template: + variable_template_parser = VariableTemplateParser(template=prompt.text) + variable_selectors.extend(variable_template_parser.extract_variable_selectors()) + elif isinstance(prompt_template, CompletionModelPromptTemplate): + variable_template_parser = VariableTemplateParser(template=prompt_template.text) + variable_selectors = variable_template_parser.extract_variable_selectors() + + for variable_selector in variable_selectors: + variable = self.graph_runtime_state.variable_pool.get(variable_selector.value_selector) + if variable is None: + raise VariableNotFoundError(f"Variable {variable_selector.variable} not found") + if isinstance(variable, NoneSegment): + inputs[variable_selector.variable] = "" + inputs[variable_selector.variable] = variable.to_object() + + memory = node_data.memory + if memory and memory.query_prompt_template: + query_variable_selectors = VariableTemplateParser( + template=memory.query_prompt_template + ).extract_variable_selectors() + for variable_selector in query_variable_selectors: + variable = self.graph_runtime_state.variable_pool.get(variable_selector.value_selector) + if variable is None: + raise VariableNotFoundError(f"Variable {variable_selector.variable} not found") + if isinstance(variable, NoneSegment): + continue + inputs[variable_selector.variable] = variable.to_object() + + return inputs + + def _fetch_files(self, *, selector: Sequence[str]) -> Sequence["File"]: + variable = self.graph_runtime_state.variable_pool.get(selector) + if variable is None: + return [] + elif isinstance(variable, FileSegment): + return [variable.value] + elif isinstance(variable, ArrayFileSegment): + return variable.value + elif isinstance(variable, NoneSegment | ArrayAnySegment): + return [] + raise InvalidVariableTypeError(f"Invalid variable type: {type(variable)}") + + def _fetch_context(self, node_data: LLMNodeData): + if not node_data.context.enabled: + return + + if not node_data.context.variable_selector: + return + + context_value_variable = self.graph_runtime_state.variable_pool.get(node_data.context.variable_selector) + if context_value_variable: + if isinstance(context_value_variable, StringSegment): + yield RunRetrieverResourceEvent(retriever_resources=[], context=context_value_variable.value) + elif isinstance(context_value_variable, ArraySegment): + context_str = "" + original_retriever_resource = [] + for item in context_value_variable.value: + if isinstance(item, str): + context_str += item + "\n" + else: + if "content" not in item: + raise InvalidContextStructureError(f"Invalid context structure: {item}") + + context_str += item["content"] + "\n" + + retriever_resource = self._convert_to_original_retriever_resource(item) + if retriever_resource: + original_retriever_resource.append(retriever_resource) + + yield RunRetrieverResourceEvent( + retriever_resources=original_retriever_resource, context=context_str.strip() + ) + + def _convert_to_original_retriever_resource(self, context_dict: dict) -> Optional[dict]: + if ( + "metadata" in context_dict + and "_source" in context_dict["metadata"] + and context_dict["metadata"]["_source"] == "knowledge" + ): + metadata = context_dict.get("metadata", {}) + + source = { + "position": metadata.get("position"), + "dataset_id": metadata.get("dataset_id"), + "dataset_name": metadata.get("dataset_name"), + "document_id": metadata.get("document_id"), + "document_name": metadata.get("document_name"), + "data_source_type": metadata.get("document_data_source_type"), + "segment_id": metadata.get("segment_id"), + "retriever_from": metadata.get("retriever_from"), + "score": metadata.get("score"), + "hit_count": metadata.get("segment_hit_count"), + "word_count": metadata.get("segment_word_count"), + "segment_position": metadata.get("segment_position"), + "index_node_hash": metadata.get("segment_index_node_hash"), + "content": context_dict.get("content"), + "page": metadata.get("page"), + } + + return source + + return None + + def _fetch_model_config( + self, node_data_model: ModelConfig + ) -> tuple[ModelInstance, ModelConfigWithCredentialsEntity]: + model_name = node_data_model.name + provider_name = node_data_model.provider + + model_manager = ModelManager() + model_instance = model_manager.get_model_instance( + tenant_id=self.tenant_id, model_type=ModelType.LLM, provider=provider_name, model=model_name + ) + + provider_model_bundle = model_instance.provider_model_bundle + model_type_instance = model_instance.model_type_instance + model_type_instance = cast(LargeLanguageModel, model_type_instance) + + model_credentials = model_instance.credentials + + # check model + provider_model = provider_model_bundle.configuration.get_provider_model( + model=model_name, model_type=ModelType.LLM + ) + + if provider_model is None: + raise ModelNotExistError(f"Model {model_name} not exist.") + + if provider_model.status == ModelStatus.NO_CONFIGURE: + raise ProviderTokenNotInitError(f"Model {model_name} credentials is not initialized.") + elif provider_model.status == ModelStatus.NO_PERMISSION: + raise ModelCurrentlyNotSupportError(f"Dify Hosted OpenAI {model_name} currently not support.") + elif provider_model.status == ModelStatus.QUOTA_EXCEEDED: + raise QuotaExceededError(f"Model provider {provider_name} quota exceeded.") + + # model config + completion_params = node_data_model.completion_params + stop = [] + if "stop" in completion_params: + stop = completion_params["stop"] + del completion_params["stop"] + + # get model mode + model_mode = node_data_model.mode + if not model_mode: + raise LLMModeRequiredError("LLM mode is required.") + + model_schema = model_type_instance.get_model_schema(model_name, model_credentials) + + if not model_schema: + raise ModelNotExistError(f"Model {model_name} not exist.") + + return model_instance, ModelConfigWithCredentialsEntity( + provider=provider_name, + model=model_name, + model_schema=model_schema, + mode=model_mode, + provider_model_bundle=provider_model_bundle, + credentials=model_credentials, + parameters=completion_params, + stop=stop, + ) + + def _fetch_memory( + self, node_data_memory: Optional[MemoryConfig], model_instance: ModelInstance + ) -> Optional[TokenBufferMemory]: + if not node_data_memory: + return None + + # get conversation id + conversation_id_variable = self.graph_runtime_state.variable_pool.get( + ["sys", SystemVariableKey.CONVERSATION_ID.value] + ) + if not isinstance(conversation_id_variable, StringSegment): + return None + conversation_id = conversation_id_variable.value + + # get conversation + conversation = ( + db.session.query(Conversation) + .filter(Conversation.app_id == self.app_id, Conversation.id == conversation_id) + .first() + ) + + if not conversation: + return None + + memory = TokenBufferMemory(conversation=conversation, model_instance=model_instance) + + return memory + + def _fetch_prompt_messages( + self, + *, + sys_query: str | None = None, + sys_files: Sequence["File"], + context: str | None = None, + memory: TokenBufferMemory | None = None, + model_config: ModelConfigWithCredentialsEntity, + prompt_template: Sequence[LLMNodeChatModelMessage] | LLMNodeCompletionModelPromptTemplate, + memory_config: MemoryConfig | None = None, + vision_enabled: bool = False, + vision_detail: ImagePromptMessageContent.DETAIL, + variable_pool: VariablePool, + jinja2_variables: Sequence[VariableSelector], + ) -> tuple[Sequence[PromptMessage], Optional[Sequence[str]]]: + # FIXME: fix the type error cause prompt_messages is type quick a few times + prompt_messages: list[Any] = [] + + if isinstance(prompt_template, list): + # For chat model + prompt_messages.extend( + self._handle_list_messages( + messages=prompt_template, + context=context, + jinja2_variables=jinja2_variables, + variable_pool=variable_pool, + vision_detail_config=vision_detail, + ) + ) + + # Get memory messages for chat mode + memory_messages = _handle_memory_chat_mode( + memory=memory, + memory_config=memory_config, + model_config=model_config, + ) + # Extend prompt_messages with memory messages + prompt_messages.extend(memory_messages) + + # Add current query to the prompt messages + if sys_query: + message = LLMNodeChatModelMessage( + text=sys_query, + role=PromptMessageRole.USER, + edition_type="basic", + ) + prompt_messages.extend( + self._handle_list_messages( + messages=[message], + context="", + jinja2_variables=[], + variable_pool=variable_pool, + vision_detail_config=vision_detail, + ) + ) + + elif isinstance(prompt_template, LLMNodeCompletionModelPromptTemplate): + # For completion model + prompt_messages.extend( + _handle_completion_template( + template=prompt_template, + context=context, + jinja2_variables=jinja2_variables, + variable_pool=variable_pool, + ) + ) + + # Get memory text for completion model + memory_text = _handle_memory_completion_mode( + memory=memory, + memory_config=memory_config, + model_config=model_config, + ) + # Insert histories into the prompt + prompt_content = prompt_messages[0].content + # For issue #11247 - Check if prompt content is a string or a list + prompt_content_type = type(prompt_content) + if prompt_content_type == str: + if "#histories#" in prompt_content: + prompt_content = prompt_content.replace("#histories#", memory_text) + else: + prompt_content = memory_text + "\n" + prompt_content + prompt_messages[0].content = prompt_content + elif prompt_content_type == list: + for content_item in prompt_content: + if content_item.type == PromptMessageContentType.TEXT: + if "#histories#" in content_item.data: + content_item.data = content_item.data.replace("#histories#", memory_text) + else: + content_item.data = memory_text + "\n" + content_item.data + else: + raise ValueError("Invalid prompt content type") + + # Add current query to the prompt message + if sys_query: + if prompt_content_type == str: + prompt_content = prompt_messages[0].content.replace("#sys.query#", sys_query) + prompt_messages[0].content = prompt_content + elif prompt_content_type == list: + for content_item in prompt_content: + if content_item.type == PromptMessageContentType.TEXT: + content_item.data = sys_query + "\n" + content_item.data + else: + raise ValueError("Invalid prompt content type") + else: + raise TemplateTypeNotSupportError(type_name=str(type(prompt_template))) + + # The sys_files will be deprecated later + if vision_enabled and sys_files: + file_prompts = [] + for file in sys_files: + file_prompt = file_manager.to_prompt_message_content(file, image_detail_config=vision_detail) + file_prompts.append(file_prompt) + # If last prompt is a user prompt, add files into its contents, + # otherwise append a new user prompt + if ( + len(prompt_messages) > 0 + and isinstance(prompt_messages[-1], UserPromptMessage) + and isinstance(prompt_messages[-1].content, list) + ): + prompt_messages[-1] = UserPromptMessage(content=prompt_messages[-1].content + file_prompts) + else: + prompt_messages.append(UserPromptMessage(content=file_prompts)) + + # Remove empty messages and filter unsupported content + filtered_prompt_messages = [] + for prompt_message in prompt_messages: + if isinstance(prompt_message.content, list): + prompt_message_content = [] + for content_item in prompt_message.content: + # Skip content if features are not defined + if not model_config.model_schema.features: + if content_item.type != PromptMessageContentType.TEXT: + continue + prompt_message_content.append(content_item) + continue + + # Skip content if corresponding feature is not supported + if ( + ( + content_item.type == PromptMessageContentType.IMAGE + and ModelFeature.VISION not in model_config.model_schema.features + ) + or ( + content_item.type == PromptMessageContentType.DOCUMENT + and ModelFeature.DOCUMENT not in model_config.model_schema.features + ) + or ( + content_item.type == PromptMessageContentType.VIDEO + and ModelFeature.VIDEO not in model_config.model_schema.features + ) + or ( + content_item.type == PromptMessageContentType.AUDIO + and ModelFeature.AUDIO not in model_config.model_schema.features + ) + ): + continue + prompt_message_content.append(content_item) + if len(prompt_message_content) == 1 and prompt_message_content[0].type == PromptMessageContentType.TEXT: + prompt_message.content = prompt_message_content[0].data + else: + prompt_message.content = prompt_message_content + if prompt_message.is_empty(): + continue + filtered_prompt_messages.append(prompt_message) + + if len(filtered_prompt_messages) == 0: + raise NoPromptFoundError( + "No prompt found in the LLM configuration. " + "Please ensure a prompt is properly configured before proceeding." + ) + + stop = model_config.stop + return filtered_prompt_messages, stop + + @classmethod + def deduct_llm_quota(cls, tenant_id: str, model_instance: ModelInstance, usage: LLMUsage) -> None: + provider_model_bundle = model_instance.provider_model_bundle + provider_configuration = provider_model_bundle.configuration + + if provider_configuration.using_provider_type != ProviderType.SYSTEM: + return + + system_configuration = provider_configuration.system_configuration + + quota_unit = None + for quota_configuration in system_configuration.quota_configurations: + if quota_configuration.quota_type == system_configuration.current_quota_type: + quota_unit = quota_configuration.quota_unit + + if quota_configuration.quota_limit == -1: + return + + break + + used_quota = None + if quota_unit: + if quota_unit == QuotaUnit.TOKENS: + used_quota = usage.total_tokens + elif quota_unit == QuotaUnit.CREDITS: + used_quota = dify_config.get_model_credits(model_instance.model) + else: + used_quota = 1 + + if used_quota is not None and system_configuration.current_quota_type is not None: + db.session.query(Provider).filter( + Provider.tenant_id == tenant_id, + Provider.provider_name == model_instance.provider, + Provider.provider_type == ProviderType.SYSTEM.value, + Provider.quota_type == system_configuration.current_quota_type.value, + Provider.quota_limit > Provider.quota_used, + ).update({"quota_used": Provider.quota_used + used_quota}) + db.session.commit() + + @classmethod + def _extract_variable_selector_to_variable_mapping( + cls, + *, + graph_config: Mapping[str, Any], + node_id: str, + node_data: LLMNodeData, + ) -> Mapping[str, Sequence[str]]: + prompt_template = node_data.prompt_template + + variable_selectors = [] + if isinstance(prompt_template, list) and all( + isinstance(prompt, LLMNodeChatModelMessage) for prompt in prompt_template + ): + for prompt in prompt_template: + if prompt.edition_type != "jinja2": + variable_template_parser = VariableTemplateParser(template=prompt.text) + variable_selectors.extend(variable_template_parser.extract_variable_selectors()) + elif isinstance(prompt_template, LLMNodeCompletionModelPromptTemplate): + if prompt_template.edition_type != "jinja2": + variable_template_parser = VariableTemplateParser(template=prompt_template.text) + variable_selectors = variable_template_parser.extract_variable_selectors() + else: + raise InvalidVariableTypeError(f"Invalid prompt template type: {type(prompt_template)}") + + variable_mapping: dict[str, Any] = {} + for variable_selector in variable_selectors: + variable_mapping[variable_selector.variable] = variable_selector.value_selector + + memory = node_data.memory + if memory and memory.query_prompt_template: + query_variable_selectors = VariableTemplateParser( + template=memory.query_prompt_template + ).extract_variable_selectors() + for variable_selector in query_variable_selectors: + variable_mapping[variable_selector.variable] = variable_selector.value_selector + + if node_data.context.enabled: + variable_mapping["#context#"] = node_data.context.variable_selector + + if node_data.vision.enabled: + variable_mapping["#files#"] = ["sys", SystemVariableKey.FILES.value] + + if node_data.memory: + variable_mapping["#sys.query#"] = ["sys", SystemVariableKey.QUERY.value] + + if node_data.prompt_config: + enable_jinja = False + + if isinstance(prompt_template, list): + for prompt in prompt_template: + if prompt.edition_type == "jinja2": + enable_jinja = True + break + else: + if prompt_template.edition_type == "jinja2": + enable_jinja = True + + if enable_jinja: + for variable_selector in node_data.prompt_config.jinja2_variables or []: + variable_mapping[variable_selector.variable] = variable_selector.value_selector + + variable_mapping = {node_id + "." + key: value for key, value in variable_mapping.items()} + + return variable_mapping + + @classmethod + def get_default_config(cls, filters: Optional[dict] = None) -> dict: + return { + "type": "llm", + "config": { + "prompt_templates": { + "chat_model": { + "prompts": [ + {"role": "system", "text": "You are a helpful AI assistant.", "edition_type": "basic"} + ] + }, + "completion_model": { + "conversation_histories_role": {"user_prefix": "Human", "assistant_prefix": "Assistant"}, + "prompt": { + "text": "Here are the chat histories between human and assistant, inside " + " XML tags.\n\n\n{{" + "#histories#}}\n\n\n\nHuman: {{#sys.query#}}\n\nAssistant:", + "edition_type": "basic", + }, + "stop": ["Human:"], + }, + } + }, + } + + def _handle_list_messages( + self, + *, + messages: Sequence[LLMNodeChatModelMessage], + context: Optional[str], + jinja2_variables: Sequence[VariableSelector], + variable_pool: VariablePool, + vision_detail_config: ImagePromptMessageContent.DETAIL, + ) -> Sequence[PromptMessage]: + prompt_messages: list[PromptMessage] = [] + for message in messages: + if message.edition_type == "jinja2": + result_text = _render_jinja2_message( + template=message.jinja2_text or "", + jinjia2_variables=jinja2_variables, + variable_pool=variable_pool, + ) + prompt_message = _combine_message_content_with_role( + contents=[TextPromptMessageContent(data=result_text)], role=message.role + ) + prompt_messages.append(prompt_message) + else: + # Get segment group from basic message + if context: + template = message.text.replace("{#context#}", context) + else: + template = message.text + segment_group = variable_pool.convert_template(template) + + # Process segments for images + file_contents = [] + for segment in segment_group.value: + if isinstance(segment, ArrayFileSegment): + for file in segment.value: + if file.type in {FileType.IMAGE, FileType.VIDEO, FileType.AUDIO, FileType.DOCUMENT}: + file_content = file_manager.to_prompt_message_content( + file, image_detail_config=vision_detail_config + ) + file_contents.append(file_content) + elif isinstance(segment, FileSegment): + file = segment.value + if file.type in {FileType.IMAGE, FileType.VIDEO, FileType.AUDIO, FileType.DOCUMENT}: + file_content = file_manager.to_prompt_message_content( + file, image_detail_config=vision_detail_config + ) + file_contents.append(file_content) + + # Create message with text from all segments + plain_text = segment_group.text + if plain_text: + prompt_message = _combine_message_content_with_role( + contents=[TextPromptMessageContent(data=plain_text)], role=message.role + ) + prompt_messages.append(prompt_message) + + if file_contents: + # Create message with image contents + prompt_message = _combine_message_content_with_role(contents=file_contents, role=message.role) + prompt_messages.append(prompt_message) + + return prompt_messages + + +def _combine_message_content_with_role(*, contents: Sequence[PromptMessageContent], role: PromptMessageRole): + match role: + case PromptMessageRole.USER: + return UserPromptMessage(content=contents) + case PromptMessageRole.ASSISTANT: + return AssistantPromptMessage(content=contents) + case PromptMessageRole.SYSTEM: + return SystemPromptMessage(content=contents) + raise NotImplementedError(f"Role {role} is not supported") + + +def _render_jinja2_message( + *, + template: str, + jinjia2_variables: Sequence[VariableSelector], + variable_pool: VariablePool, +): + if not template: + return "" + + jinjia2_inputs = {} + for jinja2_variable in jinjia2_variables: + variable = variable_pool.get(jinja2_variable.value_selector) + jinjia2_inputs[jinja2_variable.variable] = variable.to_object() if variable else "" + code_execute_resp = CodeExecutor.execute_workflow_code_template( + language=CodeLanguage.JINJA2, + code=template, + inputs=jinjia2_inputs, + ) + result_text = code_execute_resp["result"] + return result_text + + +def _calculate_rest_token( + *, prompt_messages: list[PromptMessage], model_config: ModelConfigWithCredentialsEntity +) -> int: + rest_tokens = 2000 + + model_context_tokens = model_config.model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE) + if model_context_tokens: + model_instance = ModelInstance( + provider_model_bundle=model_config.provider_model_bundle, model=model_config.model + ) + + curr_message_tokens = model_instance.get_llm_num_tokens(prompt_messages) + + max_tokens = 0 + for parameter_rule in model_config.model_schema.parameter_rules: + if parameter_rule.name == "max_tokens" or ( + parameter_rule.use_template and parameter_rule.use_template == "max_tokens" + ): + max_tokens = ( + model_config.parameters.get(parameter_rule.name) + or model_config.parameters.get(str(parameter_rule.use_template)) + or 0 + ) + + rest_tokens = model_context_tokens - max_tokens - curr_message_tokens + rest_tokens = max(rest_tokens, 0) + + return rest_tokens + + +def _handle_memory_chat_mode( + *, + memory: TokenBufferMemory | None, + memory_config: MemoryConfig | None, + model_config: ModelConfigWithCredentialsEntity, +) -> Sequence[PromptMessage]: + memory_messages: Sequence[PromptMessage] = [] + # Get messages from memory for chat model + if memory and memory_config: + rest_tokens = _calculate_rest_token(prompt_messages=[], model_config=model_config) + memory_messages = memory.get_history_prompt_messages( + max_token_limit=rest_tokens, + message_limit=memory_config.window.size if memory_config.window.enabled else None, + ) + return memory_messages + + +def _handle_memory_completion_mode( + *, + memory: TokenBufferMemory | None, + memory_config: MemoryConfig | None, + model_config: ModelConfigWithCredentialsEntity, +) -> str: + memory_text = "" + # Get history text from memory for completion model + if memory and memory_config: + rest_tokens = _calculate_rest_token(prompt_messages=[], model_config=model_config) + if not memory_config.role_prefix: + raise MemoryRolePrefixRequiredError("Memory role prefix is required for completion model.") + memory_text = memory.get_history_prompt_text( + max_token_limit=rest_tokens, + message_limit=memory_config.window.size if memory_config.window.enabled else None, + human_prefix=memory_config.role_prefix.user, + ai_prefix=memory_config.role_prefix.assistant, + ) + return memory_text + + +def _handle_completion_template( + *, + template: LLMNodeCompletionModelPromptTemplate, + context: Optional[str], + jinja2_variables: Sequence[VariableSelector], + variable_pool: VariablePool, +) -> Sequence[PromptMessage]: + """Handle completion template processing outside of LLMNode class. + + Args: + template: The completion model prompt template + context: Optional context string + jinja2_variables: Variables for jinja2 template rendering + variable_pool: Variable pool for template conversion + + Returns: + Sequence of prompt messages + """ + prompt_messages = [] + if template.edition_type == "jinja2": + result_text = _render_jinja2_message( + template=template.jinja2_text or "", + jinjia2_variables=jinja2_variables, + variable_pool=variable_pool, + ) + else: + if context: + template_text = template.text.replace("{#context#}", context) + else: + template_text = template.text + result_text = variable_pool.convert_template(template_text).text + prompt_message = _combine_message_content_with_role( + contents=[TextPromptMessageContent(data=result_text)], role=PromptMessageRole.USER + ) + prompt_messages.append(prompt_message) + return prompt_messages diff --git a/api/core/workflow/nodes/loop/__init__.py b/api/core/workflow/nodes/loop/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/workflow/nodes/loop/entities.py b/api/core/workflow/nodes/loop/entities.py new file mode 100644 index 0000000000000000000000000000000000000000..b7cd7a948e3f8b9a76f93c5537fd2e5645bbb1fd --- /dev/null +++ b/api/core/workflow/nodes/loop/entities.py @@ -0,0 +1,13 @@ +from core.workflow.nodes.base import BaseIterationNodeData, BaseIterationState + + +class LoopNodeData(BaseIterationNodeData): + """ + Loop Node Data. + """ + + +class LoopState(BaseIterationState): + """ + Loop State. + """ diff --git a/api/core/workflow/nodes/loop/loop_node.py b/api/core/workflow/nodes/loop/loop_node.py new file mode 100644 index 0000000000000000000000000000000000000000..a366c287c2ac5642c68cd879e27581cdb365b13d --- /dev/null +++ b/api/core/workflow/nodes/loop/loop_node.py @@ -0,0 +1,37 @@ +from typing import Any + +from core.workflow.nodes.base import BaseNode +from core.workflow.nodes.enums import NodeType +from core.workflow.nodes.loop.entities import LoopNodeData, LoopState +from core.workflow.utils.condition.entities import Condition + + +class LoopNode(BaseNode[LoopNodeData]): + """ + Loop Node. + """ + + _node_data_cls = LoopNodeData + _node_type = NodeType.LOOP + + def _run(self) -> LoopState: # type: ignore + return super()._run() # type: ignore + + @classmethod + def get_conditions(cls, node_config: dict[str, Any]) -> list[Condition]: + """ + Get conditions. + """ + node_id = node_config.get("id") + if not node_id: + return [] + + # TODO waiting for implementation + return [ + Condition( # type: ignore + variable_selector=[node_id, "index"], + comparison_operator="≤", + value_type="value_selector", + value_selector=[], + ) + ] diff --git a/api/core/workflow/nodes/node_mapping.py b/api/core/workflow/nodes/node_mapping.py new file mode 100644 index 0000000000000000000000000000000000000000..51fc5129cdd875c2f98a5c8b0e19d8d37ca7f79f --- /dev/null +++ b/api/core/workflow/nodes/node_mapping.py @@ -0,0 +1,104 @@ +from collections.abc import Mapping + +from core.workflow.nodes.answer import AnswerNode +from core.workflow.nodes.base import BaseNode +from core.workflow.nodes.code import CodeNode +from core.workflow.nodes.document_extractor import DocumentExtractorNode +from core.workflow.nodes.end import EndNode +from core.workflow.nodes.enums import NodeType +from core.workflow.nodes.http_request import HttpRequestNode +from core.workflow.nodes.if_else import IfElseNode +from core.workflow.nodes.iteration import IterationNode, IterationStartNode +from core.workflow.nodes.knowledge_retrieval import KnowledgeRetrievalNode +from core.workflow.nodes.list_operator import ListOperatorNode +from core.workflow.nodes.llm import LLMNode +from core.workflow.nodes.parameter_extractor import ParameterExtractorNode +from core.workflow.nodes.question_classifier import QuestionClassifierNode +from core.workflow.nodes.start import StartNode +from core.workflow.nodes.template_transform import TemplateTransformNode +from core.workflow.nodes.tool import ToolNode +from core.workflow.nodes.variable_aggregator import VariableAggregatorNode +from core.workflow.nodes.variable_assigner.v1 import VariableAssignerNode as VariableAssignerNodeV1 +from core.workflow.nodes.variable_assigner.v2 import VariableAssignerNode as VariableAssignerNodeV2 + +LATEST_VERSION = "latest" + +NODE_TYPE_CLASSES_MAPPING: Mapping[NodeType, Mapping[str, type[BaseNode]]] = { + NodeType.START: { + LATEST_VERSION: StartNode, + "1": StartNode, + }, + NodeType.END: { + LATEST_VERSION: EndNode, + "1": EndNode, + }, + NodeType.ANSWER: { + LATEST_VERSION: AnswerNode, + "1": AnswerNode, + }, + NodeType.LLM: { + LATEST_VERSION: LLMNode, + "1": LLMNode, + }, + NodeType.KNOWLEDGE_RETRIEVAL: { + LATEST_VERSION: KnowledgeRetrievalNode, + "1": KnowledgeRetrievalNode, + }, + NodeType.IF_ELSE: { + LATEST_VERSION: IfElseNode, + "1": IfElseNode, + }, + NodeType.CODE: { + LATEST_VERSION: CodeNode, + "1": CodeNode, + }, + NodeType.TEMPLATE_TRANSFORM: { + LATEST_VERSION: TemplateTransformNode, + "1": TemplateTransformNode, + }, + NodeType.QUESTION_CLASSIFIER: { + LATEST_VERSION: QuestionClassifierNode, + "1": QuestionClassifierNode, + }, + NodeType.HTTP_REQUEST: { + LATEST_VERSION: HttpRequestNode, + "1": HttpRequestNode, + }, + NodeType.TOOL: { + LATEST_VERSION: ToolNode, + "1": ToolNode, + }, + NodeType.VARIABLE_AGGREGATOR: { + LATEST_VERSION: VariableAggregatorNode, + "1": VariableAggregatorNode, + }, + NodeType.LEGACY_VARIABLE_AGGREGATOR: { + LATEST_VERSION: VariableAggregatorNode, + "1": VariableAggregatorNode, + }, # original name of VARIABLE_AGGREGATOR + NodeType.ITERATION: { + LATEST_VERSION: IterationNode, + "1": IterationNode, + }, + NodeType.ITERATION_START: { + LATEST_VERSION: IterationStartNode, + "1": IterationStartNode, + }, + NodeType.PARAMETER_EXTRACTOR: { + LATEST_VERSION: ParameterExtractorNode, + "1": ParameterExtractorNode, + }, + NodeType.VARIABLE_ASSIGNER: { + LATEST_VERSION: VariableAssignerNodeV2, + "1": VariableAssignerNodeV1, + "2": VariableAssignerNodeV2, + }, + NodeType.DOCUMENT_EXTRACTOR: { + LATEST_VERSION: DocumentExtractorNode, + "1": DocumentExtractorNode, + }, + NodeType.LIST_OPERATOR: { + LATEST_VERSION: ListOperatorNode, + "1": ListOperatorNode, + }, +} diff --git a/api/core/workflow/nodes/parameter_extractor/__init__.py b/api/core/workflow/nodes/parameter_extractor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..bdbf19a7d36d7ec98c401a25e7958fcdd8440661 --- /dev/null +++ b/api/core/workflow/nodes/parameter_extractor/__init__.py @@ -0,0 +1,3 @@ +from .parameter_extractor_node import ParameterExtractorNode + +__all__ = ["ParameterExtractorNode"] diff --git a/api/core/workflow/nodes/parameter_extractor/entities.py b/api/core/workflow/nodes/parameter_extractor/entities.py new file mode 100644 index 0000000000000000000000000000000000000000..369eb13b04e8c488edca89bbd6491c81b1354cfb --- /dev/null +++ b/api/core/workflow/nodes/parameter_extractor/entities.py @@ -0,0 +1,77 @@ +from typing import Any, Literal, Optional + +from pydantic import BaseModel, Field, field_validator + +from core.prompt.entities.advanced_prompt_entities import MemoryConfig +from core.workflow.nodes.base import BaseNodeData +from core.workflow.nodes.llm import ModelConfig, VisionConfig + + +class ParameterConfig(BaseModel): + """ + Parameter Config. + """ + + name: str + type: Literal["string", "number", "bool", "select", "array[string]", "array[number]", "array[object]"] + options: Optional[list[str]] = None + description: str + required: bool + + @field_validator("name", mode="before") + @classmethod + def validate_name(cls, value) -> str: + if not value: + raise ValueError("Parameter name is required") + if value in {"__reason", "__is_success"}: + raise ValueError("Invalid parameter name, __reason and __is_success are reserved") + return str(value) + + +class ParameterExtractorNodeData(BaseNodeData): + """ + Parameter Extractor Node Data. + """ + + model: ModelConfig + query: list[str] + parameters: list[ParameterConfig] + instruction: Optional[str] = None + memory: Optional[MemoryConfig] = None + reasoning_mode: Literal["function_call", "prompt"] + vision: VisionConfig = Field(default_factory=VisionConfig) + + @field_validator("reasoning_mode", mode="before") + @classmethod + def set_reasoning_mode(cls, v) -> str: + return v or "function_call" + + def get_parameter_json_schema(self) -> dict: + """ + Get parameter json schema. + + :return: parameter json schema + """ + parameters: dict[str, Any] = {"type": "object", "properties": {}, "required": []} + + for parameter in self.parameters: + parameter_schema: dict[str, Any] = {"description": parameter.description} + + if parameter.type in {"string", "select"}: + parameter_schema["type"] = "string" + elif parameter.type.startswith("array"): + parameter_schema["type"] = "array" + nested_type = parameter.type[6:-1] + parameter_schema["items"] = {"type": nested_type} + else: + parameter_schema["type"] = parameter.type + + if parameter.type == "select": + parameter_schema["enum"] = parameter.options + + parameters["properties"][parameter.name] = parameter_schema + + if parameter.required: + parameters["required"].append(parameter.name) + + return parameters diff --git a/api/core/workflow/nodes/parameter_extractor/exc.py b/api/core/workflow/nodes/parameter_extractor/exc.py new file mode 100644 index 0000000000000000000000000000000000000000..6511aba18569990957b8111fb45f0a12c2de67e0 --- /dev/null +++ b/api/core/workflow/nodes/parameter_extractor/exc.py @@ -0,0 +1,50 @@ +class ParameterExtractorNodeError(ValueError): + """Base error for ParameterExtractorNode.""" + + +class InvalidModelTypeError(ParameterExtractorNodeError): + """Raised when the model is not a Large Language Model.""" + + +class ModelSchemaNotFoundError(ParameterExtractorNodeError): + """Raised when the model schema is not found.""" + + +class InvalidInvokeResultError(ParameterExtractorNodeError): + """Raised when the invoke result is invalid.""" + + +class InvalidTextContentTypeError(ParameterExtractorNodeError): + """Raised when the text content type is invalid.""" + + +class InvalidNumberOfParametersError(ParameterExtractorNodeError): + """Raised when the number of parameters is invalid.""" + + +class RequiredParameterMissingError(ParameterExtractorNodeError): + """Raised when a required parameter is missing.""" + + +class InvalidSelectValueError(ParameterExtractorNodeError): + """Raised when a select value is invalid.""" + + +class InvalidNumberValueError(ParameterExtractorNodeError): + """Raised when a number value is invalid.""" + + +class InvalidBoolValueError(ParameterExtractorNodeError): + """Raised when a bool value is invalid.""" + + +class InvalidStringValueError(ParameterExtractorNodeError): + """Raised when a string value is invalid.""" + + +class InvalidArrayValueError(ParameterExtractorNodeError): + """Raised when an array value is invalid.""" + + +class InvalidModelModeError(ParameterExtractorNodeError): + """Raised when the model mode is invalid.""" diff --git a/api/core/workflow/nodes/parameter_extractor/parameter_extractor_node.py b/api/core/workflow/nodes/parameter_extractor/parameter_extractor_node.py new file mode 100644 index 0000000000000000000000000000000000000000..9c88047f2c8e573766b76de6ed5dd68491215bd2 --- /dev/null +++ b/api/core/workflow/nodes/parameter_extractor/parameter_extractor_node.py @@ -0,0 +1,799 @@ +import json +import uuid +from collections.abc import Mapping, Sequence +from typing import Any, Optional, cast + +from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity +from core.file import File +from core.memory.token_buffer_memory import TokenBufferMemory +from core.model_manager import ModelInstance +from core.model_runtime.entities.llm_entities import LLMResult, LLMUsage +from core.model_runtime.entities.message_entities import ( + AssistantPromptMessage, + PromptMessage, + PromptMessageRole, + PromptMessageTool, + ToolPromptMessage, + UserPromptMessage, +) +from core.model_runtime.entities.model_entities import ModelFeature, ModelPropertyKey +from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel +from core.model_runtime.utils.encoders import jsonable_encoder +from core.prompt.advanced_prompt_transform import AdvancedPromptTransform +from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate +from core.prompt.simple_prompt_transform import ModelMode +from core.prompt.utils.prompt_message_util import PromptMessageUtil +from core.workflow.entities.node_entities import NodeRunMetadataKey, NodeRunResult +from core.workflow.entities.variable_pool import VariablePool +from core.workflow.nodes.enums import NodeType +from core.workflow.nodes.llm import LLMNode, ModelConfig +from core.workflow.utils import variable_template_parser +from extensions.ext_database import db +from models.workflow import WorkflowNodeExecutionStatus + +from .entities import ParameterExtractorNodeData +from .exc import ( + InvalidArrayValueError, + InvalidBoolValueError, + InvalidInvokeResultError, + InvalidModelModeError, + InvalidModelTypeError, + InvalidNumberOfParametersError, + InvalidNumberValueError, + InvalidSelectValueError, + InvalidStringValueError, + InvalidTextContentTypeError, + ModelSchemaNotFoundError, + ParameterExtractorNodeError, + RequiredParameterMissingError, +) +from .prompts import ( + CHAT_EXAMPLE, + CHAT_GENERATE_JSON_USER_MESSAGE_TEMPLATE, + COMPLETION_GENERATE_JSON_PROMPT, + FUNCTION_CALLING_EXTRACTOR_EXAMPLE, + FUNCTION_CALLING_EXTRACTOR_NAME, + FUNCTION_CALLING_EXTRACTOR_SYSTEM_PROMPT, + FUNCTION_CALLING_EXTRACTOR_USER_TEMPLATE, +) + + +class ParameterExtractorNode(LLMNode): + """ + Parameter Extractor Node. + """ + + # FIXME: figure out why here is different from super class + _node_data_cls = ParameterExtractorNodeData # type: ignore + _node_type = NodeType.PARAMETER_EXTRACTOR + + _model_instance: Optional[ModelInstance] = None + _model_config: Optional[ModelConfigWithCredentialsEntity] = None + + @classmethod + def get_default_config(cls, filters: Optional[dict] = None) -> dict: + return { + "model": { + "prompt_templates": { + "completion_model": { + "conversation_histories_role": {"user_prefix": "Human", "assistant_prefix": "Assistant"}, + "stop": ["Human:"], + } + } + } + } + + def _run(self): + """ + Run the node. + """ + node_data = cast(ParameterExtractorNodeData, self.node_data) + variable = self.graph_runtime_state.variable_pool.get(node_data.query) + query = variable.text if variable else "" + + files = ( + self._fetch_files( + selector=node_data.vision.configs.variable_selector, + ) + if node_data.vision.enabled + else [] + ) + + model_instance, model_config = self._fetch_model_config(node_data.model) + if not isinstance(model_instance.model_type_instance, LargeLanguageModel): + raise InvalidModelTypeError("Model is not a Large Language Model") + + llm_model = model_instance.model_type_instance + model_schema = llm_model.get_model_schema( + model=model_config.model, + credentials=model_config.credentials, + ) + if not model_schema: + raise ModelSchemaNotFoundError("Model schema not found") + + # fetch memory + memory = self._fetch_memory( + node_data_memory=node_data.memory, + model_instance=model_instance, + ) + + if ( + set(model_schema.features or []) & {ModelFeature.TOOL_CALL, ModelFeature.MULTI_TOOL_CALL} + and node_data.reasoning_mode == "function_call" + ): + # use function call + prompt_messages, prompt_message_tools = self._generate_function_call_prompt( + node_data=node_data, + query=query, + variable_pool=self.graph_runtime_state.variable_pool, + model_config=model_config, + memory=memory, + files=files, + ) + else: + # use prompt engineering + prompt_messages = self._generate_prompt_engineering_prompt( + data=node_data, + query=query, + variable_pool=self.graph_runtime_state.variable_pool, + model_config=model_config, + memory=memory, + files=files, + ) + + prompt_message_tools = [] + + inputs = { + "query": query, + "files": [f.to_dict() for f in files], + "parameters": jsonable_encoder(node_data.parameters), + "instruction": jsonable_encoder(node_data.instruction), + } + + process_data = { + "model_mode": model_config.mode, + "prompts": PromptMessageUtil.prompt_messages_to_prompt_for_saving( + model_mode=model_config.mode, prompt_messages=prompt_messages + ), + "usage": None, + "function": {} if not prompt_message_tools else jsonable_encoder(prompt_message_tools[0]), + "tool_call": None, + } + + try: + text, usage, tool_call = self._invoke( + node_data_model=node_data.model, + model_instance=model_instance, + prompt_messages=prompt_messages, + tools=prompt_message_tools, + stop=model_config.stop, + ) + process_data["usage"] = jsonable_encoder(usage) + process_data["tool_call"] = jsonable_encoder(tool_call) + process_data["llm_text"] = text + except ParameterExtractorNodeError as e: + return NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + inputs=inputs, + process_data=process_data, + outputs={"__is_success": 0, "__reason": str(e)}, + error=str(e), + metadata={}, + ) + except Exception as e: + return NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + inputs=inputs, + process_data=process_data, + outputs={"__is_success": 0, "__reason": "Failed to invoke model", "__error": str(e)}, + error=str(e), + metadata={}, + ) + + error = None + + if tool_call: + result = self._extract_json_from_tool_call(tool_call) + else: + result = self._extract_complete_json_response(text) + if not result: + result = self._generate_default_result(node_data) + error = "Failed to extract result from function call or text response, using empty result." + + try: + result = self._validate_result(data=node_data, result=result or {}) + except ParameterExtractorNodeError as e: + error = str(e) + + # transform result into standard format + result = self._transform_result(data=node_data, result=result or {}) + + return NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + inputs=inputs, + process_data=process_data, + outputs={"__is_success": 1 if not error else 0, "__reason": error, **result}, + metadata={ + NodeRunMetadataKey.TOTAL_TOKENS: usage.total_tokens, + NodeRunMetadataKey.TOTAL_PRICE: usage.total_price, + NodeRunMetadataKey.CURRENCY: usage.currency, + }, + llm_usage=usage, + ) + + def _invoke( + self, + node_data_model: ModelConfig, + model_instance: ModelInstance, + prompt_messages: list[PromptMessage], + tools: list[PromptMessageTool], + stop: list[str], + ) -> tuple[str, LLMUsage, Optional[AssistantPromptMessage.ToolCall]]: + db.session.close() + + invoke_result = model_instance.invoke_llm( + prompt_messages=prompt_messages, + model_parameters=node_data_model.completion_params, + tools=tools, + stop=stop, + stream=False, + user=self.user_id, + ) + + # handle invoke result + if not isinstance(invoke_result, LLMResult): + raise InvalidInvokeResultError(f"Invalid invoke result: {invoke_result}") + + text = invoke_result.message.content + if not isinstance(text, str | None): + raise InvalidTextContentTypeError(f"Invalid text content type: {type(text)}. Expected str.") + + usage = invoke_result.usage + tool_call = invoke_result.message.tool_calls[0] if invoke_result.message.tool_calls else None + + # deduct quota + self.deduct_llm_quota(tenant_id=self.tenant_id, model_instance=model_instance, usage=usage) + + if text is None: + text = "" + + return text, usage, tool_call + + def _generate_function_call_prompt( + self, + node_data: ParameterExtractorNodeData, + query: str, + variable_pool: VariablePool, + model_config: ModelConfigWithCredentialsEntity, + memory: Optional[TokenBufferMemory], + files: Sequence[File], + ) -> tuple[list[PromptMessage], list[PromptMessageTool]]: + """ + Generate function call prompt. + """ + query = FUNCTION_CALLING_EXTRACTOR_USER_TEMPLATE.format( + content=query, structure=json.dumps(node_data.get_parameter_json_schema()) + ) + + prompt_transform = AdvancedPromptTransform(with_variable_tmpl=True) + rest_token = self._calculate_rest_token(node_data, query, variable_pool, model_config, "") + prompt_template = self._get_function_calling_prompt_template( + node_data, query, variable_pool, memory, rest_token + ) + prompt_messages = prompt_transform.get_prompt( + prompt_template=prompt_template, + inputs={}, + query="", + files=files, + context="", + memory_config=node_data.memory, + memory=None, + model_config=model_config, + ) + + # find last user message + last_user_message_idx = -1 + for i, prompt_message in enumerate(prompt_messages): + if prompt_message.role == PromptMessageRole.USER: + last_user_message_idx = i + + # add function call messages before last user message + example_messages = [] + for example in FUNCTION_CALLING_EXTRACTOR_EXAMPLE: + id = uuid.uuid4().hex + example_messages.extend( + [ + UserPromptMessage(content=example["user"]["query"]), + AssistantPromptMessage( + content=example["assistant"]["text"], + tool_calls=[ + AssistantPromptMessage.ToolCall( + id=id, + type="function", + function=AssistantPromptMessage.ToolCall.ToolCallFunction( + name=example["assistant"]["function_call"]["name"], + arguments=json.dumps(example["assistant"]["function_call"]["parameters"]), + ), + ) + ], + ), + ToolPromptMessage( + content="Great! You have called the function with the correct parameters.", tool_call_id=id + ), + AssistantPromptMessage( + content="I have extracted the parameters, let's move on.", + ), + ] + ) + + prompt_messages = ( + prompt_messages[:last_user_message_idx] + example_messages + prompt_messages[last_user_message_idx:] + ) + + # generate tool + tool = PromptMessageTool( + name=FUNCTION_CALLING_EXTRACTOR_NAME, + description="Extract parameters from the natural language text", + parameters=node_data.get_parameter_json_schema(), + ) + + return prompt_messages, [tool] + + def _generate_prompt_engineering_prompt( + self, + data: ParameterExtractorNodeData, + query: str, + variable_pool: VariablePool, + model_config: ModelConfigWithCredentialsEntity, + memory: Optional[TokenBufferMemory], + files: Sequence[File], + ) -> list[PromptMessage]: + """ + Generate prompt engineering prompt. + """ + model_mode = ModelMode.value_of(data.model.mode) + + if model_mode == ModelMode.COMPLETION: + return self._generate_prompt_engineering_completion_prompt( + node_data=data, + query=query, + variable_pool=variable_pool, + model_config=model_config, + memory=memory, + files=files, + ) + elif model_mode == ModelMode.CHAT: + return self._generate_prompt_engineering_chat_prompt( + node_data=data, + query=query, + variable_pool=variable_pool, + model_config=model_config, + memory=memory, + files=files, + ) + else: + raise InvalidModelModeError(f"Invalid model mode: {model_mode}") + + def _generate_prompt_engineering_completion_prompt( + self, + node_data: ParameterExtractorNodeData, + query: str, + variable_pool: VariablePool, + model_config: ModelConfigWithCredentialsEntity, + memory: Optional[TokenBufferMemory], + files: Sequence[File], + ) -> list[PromptMessage]: + """ + Generate completion prompt. + """ + prompt_transform = AdvancedPromptTransform(with_variable_tmpl=True) + rest_token = self._calculate_rest_token( + node_data=node_data, query=query, variable_pool=variable_pool, model_config=model_config, context="" + ) + prompt_template = self._get_prompt_engineering_prompt_template( + node_data=node_data, query=query, variable_pool=variable_pool, memory=memory, max_token_limit=rest_token + ) + prompt_messages = prompt_transform.get_prompt( + prompt_template=prompt_template, + inputs={"structure": json.dumps(node_data.get_parameter_json_schema())}, + query="", + files=files, + context="", + memory_config=node_data.memory, + memory=memory, + model_config=model_config, + ) + + return prompt_messages + + def _generate_prompt_engineering_chat_prompt( + self, + node_data: ParameterExtractorNodeData, + query: str, + variable_pool: VariablePool, + model_config: ModelConfigWithCredentialsEntity, + memory: Optional[TokenBufferMemory], + files: Sequence[File], + ) -> list[PromptMessage]: + """ + Generate chat prompt. + """ + prompt_transform = AdvancedPromptTransform(with_variable_tmpl=True) + rest_token = self._calculate_rest_token( + node_data=node_data, query=query, variable_pool=variable_pool, model_config=model_config, context="" + ) + prompt_template = self._get_prompt_engineering_prompt_template( + node_data=node_data, + query=CHAT_GENERATE_JSON_USER_MESSAGE_TEMPLATE.format( + structure=json.dumps(node_data.get_parameter_json_schema()), text=query + ), + variable_pool=variable_pool, + memory=memory, + max_token_limit=rest_token, + ) + + prompt_messages = prompt_transform.get_prompt( + prompt_template=prompt_template, + inputs={}, + query="", + files=files, + context="", + memory_config=node_data.memory, + memory=None, + model_config=model_config, + ) + + # find last user message + last_user_message_idx = -1 + for i, prompt_message in enumerate(prompt_messages): + if prompt_message.role == PromptMessageRole.USER: + last_user_message_idx = i + + # add example messages before last user message + example_messages = [] + for example in CHAT_EXAMPLE: + example_messages.extend( + [ + UserPromptMessage( + content=CHAT_GENERATE_JSON_USER_MESSAGE_TEMPLATE.format( + structure=json.dumps(example["user"]["json"]), + text=example["user"]["query"], + ) + ), + AssistantPromptMessage( + content=json.dumps(example["assistant"]["json"]), + ), + ] + ) + + prompt_messages = ( + prompt_messages[:last_user_message_idx] + example_messages + prompt_messages[last_user_message_idx:] + ) + + return prompt_messages + + def _validate_result(self, data: ParameterExtractorNodeData, result: dict) -> dict: + """ + Validate result. + """ + if len(data.parameters) != len(result): + raise InvalidNumberOfParametersError("Invalid number of parameters") + + for parameter in data.parameters: + if parameter.required and parameter.name not in result: + raise RequiredParameterMissingError(f"Parameter {parameter.name} is required") + + if parameter.type == "select" and parameter.options and result.get(parameter.name) not in parameter.options: + raise InvalidSelectValueError(f"Invalid `select` value for parameter {parameter.name}") + + if parameter.type == "number" and not isinstance(result.get(parameter.name), int | float): + raise InvalidNumberValueError(f"Invalid `number` value for parameter {parameter.name}") + + if parameter.type == "bool" and not isinstance(result.get(parameter.name), bool): + raise InvalidBoolValueError(f"Invalid `bool` value for parameter {parameter.name}") + + if parameter.type == "string" and not isinstance(result.get(parameter.name), str): + raise InvalidStringValueError(f"Invalid `string` value for parameter {parameter.name}") + + if parameter.type.startswith("array"): + parameters = result.get(parameter.name) + if not isinstance(parameters, list): + raise InvalidArrayValueError(f"Invalid `array` value for parameter {parameter.name}") + nested_type = parameter.type[6:-1] + for item in parameters: + if nested_type == "number" and not isinstance(item, int | float): + raise InvalidArrayValueError(f"Invalid `array[number]` value for parameter {parameter.name}") + if nested_type == "string" and not isinstance(item, str): + raise InvalidArrayValueError(f"Invalid `array[string]` value for parameter {parameter.name}") + if nested_type == "object" and not isinstance(item, dict): + raise InvalidArrayValueError(f"Invalid `array[object]` value for parameter {parameter.name}") + return result + + def _transform_result(self, data: ParameterExtractorNodeData, result: dict) -> dict: + """ + Transform result into standard format. + """ + transformed_result = {} + for parameter in data.parameters: + if parameter.name in result: + # transform value + if parameter.type == "number": + if isinstance(result[parameter.name], int | float): + transformed_result[parameter.name] = result[parameter.name] + elif isinstance(result[parameter.name], str): + try: + if "." in result[parameter.name]: + result[parameter.name] = float(result[parameter.name]) + else: + result[parameter.name] = int(result[parameter.name]) + except ValueError: + pass + else: + pass + # TODO: bool is not supported in the current version + # elif parameter.type == 'bool': + # if isinstance(result[parameter.name], bool): + # transformed_result[parameter.name] = bool(result[parameter.name]) + # elif isinstance(result[parameter.name], str): + # if result[parameter.name].lower() in ['true', 'false']: + # transformed_result[parameter.name] = bool(result[parameter.name].lower() == 'true') + # elif isinstance(result[parameter.name], int): + # transformed_result[parameter.name] = bool(result[parameter.name]) + elif parameter.type in {"string", "select"}: + if isinstance(result[parameter.name], str): + transformed_result[parameter.name] = result[parameter.name] + elif parameter.type.startswith("array"): + if isinstance(result[parameter.name], list): + nested_type = parameter.type[6:-1] + transformed_result[parameter.name] = [] + for item in result[parameter.name]: + if nested_type == "number": + if isinstance(item, int | float): + transformed_result[parameter.name].append(item) + elif isinstance(item, str): + try: + if "." in item: + transformed_result[parameter.name].append(float(item)) + else: + transformed_result[parameter.name].append(int(item)) + except ValueError: + pass + elif nested_type == "string": + if isinstance(item, str): + transformed_result[parameter.name].append(item) + elif nested_type == "object": + if isinstance(item, dict): + transformed_result[parameter.name].append(item) + + if parameter.name not in transformed_result: + if parameter.type == "number": + transformed_result[parameter.name] = 0 + elif parameter.type == "bool": + transformed_result[parameter.name] = False + elif parameter.type in {"string", "select"}: + transformed_result[parameter.name] = "" + elif parameter.type.startswith("array"): + transformed_result[parameter.name] = [] + + return transformed_result + + def _extract_complete_json_response(self, result: str) -> Optional[dict]: + """ + Extract complete json response. + """ + + def extract_json(text): + """ + From a given JSON started from '{' or '[' extract the complete JSON object. + """ + stack = [] + for i, c in enumerate(text): + if c in {"{", "["}: + stack.append(c) + elif c in {"}", "]"}: + # check if stack is empty + if not stack: + return text[:i] + # check if the last element in stack is matching + if (c == "}" and stack[-1] == "{") or (c == "]" and stack[-1] == "["): + stack.pop() + if not stack: + return text[: i + 1] + else: + return text[:i] + return None + + # extract json from the text + for idx in range(len(result)): + if result[idx] == "{" or result[idx] == "[": + json_str = extract_json(result[idx:]) + if json_str: + try: + return cast(dict, json.loads(json_str)) + except Exception: + pass + return None + + def _extract_json_from_tool_call(self, tool_call: AssistantPromptMessage.ToolCall) -> Optional[dict]: + """ + Extract json from tool call. + """ + if not tool_call or not tool_call.function.arguments: + return None + + return cast(dict, json.loads(tool_call.function.arguments)) + + def _generate_default_result(self, data: ParameterExtractorNodeData) -> dict: + """ + Generate default result. + """ + result: dict[str, Any] = {} + for parameter in data.parameters: + if parameter.type == "number": + result[parameter.name] = 0 + elif parameter.type == "bool": + result[parameter.name] = False + elif parameter.type in {"string", "select"}: + result[parameter.name] = "" + + return result + + def _get_function_calling_prompt_template( + self, + node_data: ParameterExtractorNodeData, + query: str, + variable_pool: VariablePool, + memory: Optional[TokenBufferMemory], + max_token_limit: int = 2000, + ) -> list[ChatModelMessage]: + model_mode = ModelMode.value_of(node_data.model.mode) + input_text = query + memory_str = "" + instruction = variable_pool.convert_template(node_data.instruction or "").text + + if memory and node_data.memory and node_data.memory.window: + memory_str = memory.get_history_prompt_text( + max_token_limit=max_token_limit, message_limit=node_data.memory.window.size + ) + if model_mode == ModelMode.CHAT: + system_prompt_messages = ChatModelMessage( + role=PromptMessageRole.SYSTEM, + text=FUNCTION_CALLING_EXTRACTOR_SYSTEM_PROMPT.format(histories=memory_str, instruction=instruction), + ) + user_prompt_message = ChatModelMessage(role=PromptMessageRole.USER, text=input_text) + return [system_prompt_messages, user_prompt_message] + else: + raise InvalidModelModeError(f"Model mode {model_mode} not support.") + + def _get_prompt_engineering_prompt_template( + self, + node_data: ParameterExtractorNodeData, + query: str, + variable_pool: VariablePool, + memory: Optional[TokenBufferMemory], + max_token_limit: int = 2000, + ): + model_mode = ModelMode.value_of(node_data.model.mode) + input_text = query + memory_str = "" + instruction = variable_pool.convert_template(node_data.instruction or "").text + + if memory and node_data.memory and node_data.memory.window: + memory_str = memory.get_history_prompt_text( + max_token_limit=max_token_limit, message_limit=node_data.memory.window.size + ) + if model_mode == ModelMode.CHAT: + system_prompt_messages = ChatModelMessage( + role=PromptMessageRole.SYSTEM, + text=FUNCTION_CALLING_EXTRACTOR_SYSTEM_PROMPT.format(histories=memory_str, instruction=instruction), + ) + user_prompt_message = ChatModelMessage(role=PromptMessageRole.USER, text=input_text) + return [system_prompt_messages, user_prompt_message] + elif model_mode == ModelMode.COMPLETION: + return CompletionModelPromptTemplate( + text=COMPLETION_GENERATE_JSON_PROMPT.format( + histories=memory_str, text=input_text, instruction=instruction + ) + .replace("{γγγ", "") + .replace("}γγγ", "") + ) + else: + raise InvalidModelModeError(f"Model mode {model_mode} not support.") + + def _calculate_rest_token( + self, + node_data: ParameterExtractorNodeData, + query: str, + variable_pool: VariablePool, + model_config: ModelConfigWithCredentialsEntity, + context: Optional[str], + ) -> int: + prompt_transform = AdvancedPromptTransform(with_variable_tmpl=True) + + model_instance, model_config = self._fetch_model_config(node_data.model) + if not isinstance(model_instance.model_type_instance, LargeLanguageModel): + raise InvalidModelTypeError("Model is not a Large Language Model") + + llm_model = model_instance.model_type_instance + model_schema = llm_model.get_model_schema(model_config.model, model_config.credentials) + if not model_schema: + raise ModelSchemaNotFoundError("Model schema not found") + + if set(model_schema.features or []) & {ModelFeature.MULTI_TOOL_CALL, ModelFeature.MULTI_TOOL_CALL}: + prompt_template = self._get_function_calling_prompt_template(node_data, query, variable_pool, None, 2000) + else: + prompt_template = self._get_prompt_engineering_prompt_template(node_data, query, variable_pool, None, 2000) + + prompt_messages = prompt_transform.get_prompt( + prompt_template=prompt_template, + inputs={}, + query="", + files=[], + context=context, + memory_config=node_data.memory, + memory=None, + model_config=model_config, + ) + rest_tokens = 2000 + + model_context_tokens = model_config.model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE) + if model_context_tokens: + model_type_instance = model_config.provider_model_bundle.model_type_instance + model_type_instance = cast(LargeLanguageModel, model_type_instance) + + curr_message_tokens = ( + model_type_instance.get_num_tokens(model_config.model, model_config.credentials, prompt_messages) + 1000 + ) # add 1000 to ensure tool call messages + + max_tokens = 0 + for parameter_rule in model_config.model_schema.parameter_rules: + if parameter_rule.name == "max_tokens" or ( + parameter_rule.use_template and parameter_rule.use_template == "max_tokens" + ): + max_tokens = ( + model_config.parameters.get(parameter_rule.name) + or model_config.parameters.get(parameter_rule.use_template or "") + ) or 0 + + rest_tokens = model_context_tokens - max_tokens - curr_message_tokens + rest_tokens = max(rest_tokens, 0) + + return rest_tokens + + def _fetch_model_config( + self, node_data_model: ModelConfig + ) -> tuple[ModelInstance, ModelConfigWithCredentialsEntity]: + """ + Fetch model config. + """ + if not self._model_instance or not self._model_config: + self._model_instance, self._model_config = super()._fetch_model_config(node_data_model) + + return self._model_instance, self._model_config + + @classmethod + def _extract_variable_selector_to_variable_mapping( + cls, + *, + graph_config: Mapping[str, Any], + node_id: str, + node_data: ParameterExtractorNodeData, # type: ignore + ) -> Mapping[str, Sequence[str]]: + """ + Extract variable selector to variable mapping + :param graph_config: graph config + :param node_id: node id + :param node_data: node data + :return: + """ + # FIXME: fix the type error later + variable_mapping: dict[str, Sequence[str]] = {"query": node_data.query} + + if node_data.instruction: + selectors = variable_template_parser.extract_selectors_from_template(node_data.instruction) + for selector in selectors: + variable_mapping[selector.variable] = selector.value_selector + + variable_mapping = {node_id + "." + key: value for key, value in variable_mapping.items()} + + return variable_mapping diff --git a/api/core/workflow/nodes/parameter_extractor/prompts.py b/api/core/workflow/nodes/parameter_extractor/prompts.py new file mode 100644 index 0000000000000000000000000000000000000000..6c3155ac9a54e3fe1816a3952e95aee7ebd62932 --- /dev/null +++ b/api/core/workflow/nodes/parameter_extractor/prompts.py @@ -0,0 +1,184 @@ +from typing import Any + +FUNCTION_CALLING_EXTRACTOR_NAME = "extract_parameters" + +FUNCTION_CALLING_EXTRACTOR_SYSTEM_PROMPT = f"""You are a helpful assistant tasked with extracting structured information based on specific criteria provided. Follow the guidelines below to ensure consistency and accuracy. +### Task +Always call the `{FUNCTION_CALLING_EXTRACTOR_NAME}` function with the correct parameters. Ensure that the information extraction is contextual and aligns with the provided criteria. +### Memory +Here is the chat history between the human and assistant, provided within tags: + +\x7bhistories\x7d + +### Instructions: +Some additional information is provided below. Always adhere to these instructions as closely as possible: + +\x7binstruction\x7d + +Steps: +1. Review the chat history provided within the tags. +2. Extract the relevant information based on the criteria given, output multiple values if there is multiple relevant information that match the criteria in the given text. +3. Generate a well-formatted output using the defined functions and arguments. +4. Use the `extract_parameter` function to create structured outputs with appropriate parameters. +5. Do not include any XML tags in your output. +### Example +To illustrate, if the task involves extracting a user's name and their request, your function call might look like this: Ensure your output follows a similar structure to examples. +### Final Output +Produce well-formatted function calls in json without XML tags, as shown in the example. +""" # noqa: E501 + +FUNCTION_CALLING_EXTRACTOR_USER_TEMPLATE = f"""extract structured information from context inside XML tags by calling the function {FUNCTION_CALLING_EXTRACTOR_NAME} with the correct parameters with structure inside XML tags. + +\x7bcontent\x7d + + + +\x7bstructure\x7d + +""" # noqa: E501 + +FUNCTION_CALLING_EXTRACTOR_EXAMPLE: list[dict[str, Any]] = [ + { + "user": { + "query": "What is the weather today in SF?", + "function": { + "name": FUNCTION_CALLING_EXTRACTOR_NAME, + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The location to get the weather information", + "required": True, + }, + }, + "required": ["location"], + }, + }, + }, + "assistant": { + "text": "I need always call the function with the correct parameters." + " in this case, I need to call the function with the location parameter.", + "function_call": {"name": FUNCTION_CALLING_EXTRACTOR_NAME, "parameters": {"location": "San Francisco"}}, + }, + }, + { + "user": { + "query": "I want to eat some apple pie.", + "function": { + "name": FUNCTION_CALLING_EXTRACTOR_NAME, + "parameters": { + "type": "object", + "properties": {"food": {"type": "string", "description": "The food to eat", "required": True}}, + "required": ["food"], + }, + }, + }, + "assistant": { + "text": "I need always call the function with the correct parameters." + " in this case, I need to call the function with the food parameter.", + "function_call": {"name": FUNCTION_CALLING_EXTRACTOR_NAME, "parameters": {"food": "apple pie"}}, + }, + }, +] + +COMPLETION_GENERATE_JSON_PROMPT = """### Instructions: +Some extra information are provided below, I should always follow the instructions as possible as I can. + +{instruction} + + +### Extract parameter Workflow +I need to extract the following information from the input text. The tag specifies the 'type', 'description' and 'required' of the information to be extracted. + +{{ structure }} + + +Step 1: Carefully read the input and understand the structure of the expected output. +Step 2: Extract relevant parameters from the provided text based on the name and description of object. +Step 3: Structure the extracted parameters to JSON object as specified in . +Step 4: Ensure that the JSON object is properly formatted and valid. The output should not contain any XML tags. Only the JSON object should be outputted. + +### Memory +Here are the chat histories between human and assistant, inside XML tags. + +{histories} + + +### Structure +Here is the structure of the expected output, I should always follow the output structure. +{{γγγ + 'properties1': 'relevant text extracted from input', + 'properties2': 'relevant text extracted from input', +}}γγγ + +### Input Text +Inside XML tags, there is a text that I should extract parameters and convert to a JSON object. + +{text} + + +### Answer +I should always output a valid JSON object. Output nothing other than the JSON object. +```JSON +""" # noqa: E501 + +CHAT_GENERATE_JSON_PROMPT = """You should always follow the instructions and output a valid JSON object. +The structure of the JSON object you can found in the instructions. + +### Memory +Here are the chat histories between human and assistant, inside XML tags. + +{histories} + + +### Instructions: +Some extra information are provided below, you should always follow the instructions as possible as you can. + +{{instructions}} + +""" + +CHAT_GENERATE_JSON_USER_MESSAGE_TEMPLATE = """### Structure +Here is the structure of the JSON object, you should always follow the structure. + +{structure} + + +### Text to be converted to JSON +Inside XML tags, there is a text that you should convert to a JSON object. + +{text} + +""" + +CHAT_EXAMPLE = [ + { + "user": { + "query": "What is the weather today in SF?", + "json": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The location to get the weather information", + "required": True, + } + }, + "required": ["location"], + }, + }, + "assistant": {"text": "I need to output a valid JSON object.", "json": {"location": "San Francisco"}}, + }, + { + "user": { + "query": "I want to eat some apple pie.", + "json": { + "type": "object", + "properties": {"food": {"type": "string", "description": "The food to eat", "required": True}}, + "required": ["food"], + }, + }, + "assistant": {"text": "I need to output a valid JSON object.", "json": {"result": "apple pie"}}, + }, +] diff --git a/api/core/workflow/nodes/question_classifier/__init__.py b/api/core/workflow/nodes/question_classifier/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4d06b6bea366c7783ead76bbde888f991713d0aa --- /dev/null +++ b/api/core/workflow/nodes/question_classifier/__init__.py @@ -0,0 +1,4 @@ +from .entities import QuestionClassifierNodeData +from .question_classifier_node import QuestionClassifierNode + +__all__ = ["QuestionClassifierNode", "QuestionClassifierNodeData"] diff --git a/api/core/workflow/nodes/question_classifier/entities.py b/api/core/workflow/nodes/question_classifier/entities.py new file mode 100644 index 0000000000000000000000000000000000000000..5219f11d267c07c341bb7a8b0d960da175b2662d --- /dev/null +++ b/api/core/workflow/nodes/question_classifier/entities.py @@ -0,0 +1,21 @@ +from typing import Optional + +from pydantic import BaseModel, Field + +from core.prompt.entities.advanced_prompt_entities import MemoryConfig +from core.workflow.nodes.base import BaseNodeData +from core.workflow.nodes.llm import ModelConfig, VisionConfig + + +class ClassConfig(BaseModel): + id: str + name: str + + +class QuestionClassifierNodeData(BaseNodeData): + query_variable_selector: list[str] + model: ModelConfig + classes: list[ClassConfig] + instruction: Optional[str] = None + memory: Optional[MemoryConfig] = None + vision: VisionConfig = Field(default_factory=VisionConfig) diff --git a/api/core/workflow/nodes/question_classifier/exc.py b/api/core/workflow/nodes/question_classifier/exc.py new file mode 100644 index 0000000000000000000000000000000000000000..2c6354e2a70237d70af13293f68e0a6bd74b54c2 --- /dev/null +++ b/api/core/workflow/nodes/question_classifier/exc.py @@ -0,0 +1,6 @@ +class QuestionClassifierNodeError(ValueError): + """Base class for QuestionClassifierNode errors.""" + + +class InvalidModelTypeError(QuestionClassifierNodeError): + """Raised when the model is not a Large Language Model.""" diff --git a/api/core/workflow/nodes/question_classifier/question_classifier_node.py b/api/core/workflow/nodes/question_classifier/question_classifier_node.py new file mode 100644 index 0000000000000000000000000000000000000000..0ec44eefacf52f78cc683414862a9e6cc8ae7cab --- /dev/null +++ b/api/core/workflow/nodes/question_classifier/question_classifier_node.py @@ -0,0 +1,308 @@ +import json +from collections.abc import Mapping, Sequence +from typing import Any, Optional, cast + +from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity +from core.memory.token_buffer_memory import TokenBufferMemory +from core.model_manager import ModelInstance +from core.model_runtime.entities import LLMUsage, ModelPropertyKey, PromptMessageRole +from core.model_runtime.utils.encoders import jsonable_encoder +from core.prompt.advanced_prompt_transform import AdvancedPromptTransform +from core.prompt.simple_prompt_transform import ModelMode +from core.prompt.utils.prompt_message_util import PromptMessageUtil +from core.workflow.entities.node_entities import NodeRunMetadataKey, NodeRunResult +from core.workflow.nodes.enums import NodeType +from core.workflow.nodes.event import ModelInvokeCompletedEvent +from core.workflow.nodes.llm import ( + LLMNode, + LLMNodeChatModelMessage, + LLMNodeCompletionModelPromptTemplate, +) +from core.workflow.utils.variable_template_parser import VariableTemplateParser +from libs.json_in_md_parser import parse_and_check_json_markdown +from models.workflow import WorkflowNodeExecutionStatus + +from .entities import QuestionClassifierNodeData +from .exc import InvalidModelTypeError +from .template_prompts import ( + QUESTION_CLASSIFIER_ASSISTANT_PROMPT_1, + QUESTION_CLASSIFIER_ASSISTANT_PROMPT_2, + QUESTION_CLASSIFIER_COMPLETION_PROMPT, + QUESTION_CLASSIFIER_SYSTEM_PROMPT, + QUESTION_CLASSIFIER_USER_PROMPT_1, + QUESTION_CLASSIFIER_USER_PROMPT_2, + QUESTION_CLASSIFIER_USER_PROMPT_3, +) + + +class QuestionClassifierNode(LLMNode): + _node_data_cls = QuestionClassifierNodeData # type: ignore + _node_type = NodeType.QUESTION_CLASSIFIER + + def _run(self): + node_data = cast(QuestionClassifierNodeData, self.node_data) + variable_pool = self.graph_runtime_state.variable_pool + + # extract variables + variable = variable_pool.get(node_data.query_variable_selector) if node_data.query_variable_selector else None + query = variable.value if variable else None + variables = {"query": query} + # fetch model config + model_instance, model_config = self._fetch_model_config(node_data.model) + # fetch memory + memory = self._fetch_memory( + node_data_memory=node_data.memory, + model_instance=model_instance, + ) + # fetch instruction + node_data.instruction = node_data.instruction or "" + node_data.instruction = variable_pool.convert_template(node_data.instruction).text + + files = ( + self._fetch_files( + selector=node_data.vision.configs.variable_selector, + ) + if node_data.vision.enabled + else [] + ) + + # fetch prompt messages + rest_token = self._calculate_rest_token( + node_data=node_data, + query=query or "", + model_config=model_config, + context="", + ) + prompt_template = self._get_prompt_template( + node_data=node_data, + query=query or "", + memory=memory, + max_token_limit=rest_token, + ) + prompt_messages, stop = self._fetch_prompt_messages( + prompt_template=prompt_template, + sys_query=query, + memory=memory, + model_config=model_config, + sys_files=files, + vision_enabled=node_data.vision.enabled, + vision_detail=node_data.vision.configs.detail, + variable_pool=variable_pool, + jinja2_variables=[], + ) + + result_text = "" + usage = LLMUsage.empty_usage() + finish_reason = None + + try: + # handle invoke result + generator = self._invoke_llm( + node_data_model=node_data.model, + model_instance=model_instance, + prompt_messages=prompt_messages, + stop=stop, + ) + + for event in generator: + if isinstance(event, ModelInvokeCompletedEvent): + result_text = event.text + usage = event.usage + finish_reason = event.finish_reason + break + + category_name = node_data.classes[0].name + category_id = node_data.classes[0].id + result_text_json = parse_and_check_json_markdown(result_text, []) + # result_text_json = json.loads(result_text.strip('```JSON\n')) + if "category_name" in result_text_json and "category_id" in result_text_json: + category_id_result = result_text_json["category_id"] + classes = node_data.classes + classes_map = {class_.id: class_.name for class_ in classes} + category_ids = [_class.id for _class in classes] + if category_id_result in category_ids: + category_name = classes_map[category_id_result] + category_id = category_id_result + process_data = { + "model_mode": model_config.mode, + "prompts": PromptMessageUtil.prompt_messages_to_prompt_for_saving( + model_mode=model_config.mode, prompt_messages=prompt_messages + ), + "usage": jsonable_encoder(usage), + "finish_reason": finish_reason, + } + outputs = {"class_name": category_name, "class_id": category_id} + + return NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + inputs=variables, + process_data=process_data, + outputs=outputs, + edge_source_handle=category_id, + metadata={ + NodeRunMetadataKey.TOTAL_TOKENS: usage.total_tokens, + NodeRunMetadataKey.TOTAL_PRICE: usage.total_price, + NodeRunMetadataKey.CURRENCY: usage.currency, + }, + llm_usage=usage, + ) + except ValueError as e: + return NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + inputs=variables, + error=str(e), + metadata={ + NodeRunMetadataKey.TOTAL_TOKENS: usage.total_tokens, + NodeRunMetadataKey.TOTAL_PRICE: usage.total_price, + NodeRunMetadataKey.CURRENCY: usage.currency, + }, + llm_usage=usage, + ) + + @classmethod + def _extract_variable_selector_to_variable_mapping( + cls, + *, + graph_config: Mapping[str, Any], + node_id: str, + node_data: Any, + ) -> Mapping[str, Sequence[str]]: + """ + Extract variable selector to variable mapping + :param graph_config: graph config + :param node_id: node id + :param node_data: node data + :return: + """ + node_data = cast(QuestionClassifierNodeData, node_data) + variable_mapping = {"query": node_data.query_variable_selector} + variable_selectors = [] + if node_data.instruction: + variable_template_parser = VariableTemplateParser(template=node_data.instruction) + variable_selectors.extend(variable_template_parser.extract_variable_selectors()) + for variable_selector in variable_selectors: + variable_mapping[variable_selector.variable] = variable_selector.value_selector + + variable_mapping = {node_id + "." + key: value for key, value in variable_mapping.items()} + + return variable_mapping + + @classmethod + def get_default_config(cls, filters: Optional[dict] = None) -> dict: + """ + Get default config of node. + :param filters: filter by node config parameters. + :return: + """ + return {"type": "question-classifier", "config": {"instructions": ""}} + + def _calculate_rest_token( + self, + node_data: QuestionClassifierNodeData, + query: str, + model_config: ModelConfigWithCredentialsEntity, + context: Optional[str], + ) -> int: + prompt_transform = AdvancedPromptTransform(with_variable_tmpl=True) + prompt_template = self._get_prompt_template(node_data, query, None, 2000) + prompt_messages = prompt_transform.get_prompt( + prompt_template=prompt_template, + inputs={}, + query="", + files=[], + context=context, + memory_config=node_data.memory, + memory=None, + model_config=model_config, + ) + rest_tokens = 2000 + + model_context_tokens = model_config.model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE) + if model_context_tokens: + model_instance = ModelInstance( + provider_model_bundle=model_config.provider_model_bundle, model=model_config.model + ) + + curr_message_tokens = model_instance.get_llm_num_tokens(prompt_messages) + + max_tokens = 0 + for parameter_rule in model_config.model_schema.parameter_rules: + if parameter_rule.name == "max_tokens" or ( + parameter_rule.use_template and parameter_rule.use_template == "max_tokens" + ): + max_tokens = ( + model_config.parameters.get(parameter_rule.name) + or model_config.parameters.get(parameter_rule.use_template or "") + ) or 0 + + rest_tokens = model_context_tokens - max_tokens - curr_message_tokens + rest_tokens = max(rest_tokens, 0) + + return rest_tokens + + def _get_prompt_template( + self, + node_data: QuestionClassifierNodeData, + query: str, + memory: Optional[TokenBufferMemory], + max_token_limit: int = 2000, + ): + model_mode = ModelMode.value_of(node_data.model.mode) + classes = node_data.classes + categories = [] + for class_ in classes: + category = {"category_id": class_.id, "category_name": class_.name} + categories.append(category) + instruction = node_data.instruction or "" + input_text = query + memory_str = "" + if memory: + memory_str = memory.get_history_prompt_text( + max_token_limit=max_token_limit, + message_limit=node_data.memory.window.size if node_data.memory and node_data.memory.window else None, + ) + prompt_messages: list[LLMNodeChatModelMessage] = [] + if model_mode == ModelMode.CHAT: + system_prompt_messages = LLMNodeChatModelMessage( + role=PromptMessageRole.SYSTEM, text=QUESTION_CLASSIFIER_SYSTEM_PROMPT.format(histories=memory_str) + ) + prompt_messages.append(system_prompt_messages) + user_prompt_message_1 = LLMNodeChatModelMessage( + role=PromptMessageRole.USER, text=QUESTION_CLASSIFIER_USER_PROMPT_1 + ) + prompt_messages.append(user_prompt_message_1) + assistant_prompt_message_1 = LLMNodeChatModelMessage( + role=PromptMessageRole.ASSISTANT, text=QUESTION_CLASSIFIER_ASSISTANT_PROMPT_1 + ) + prompt_messages.append(assistant_prompt_message_1) + user_prompt_message_2 = LLMNodeChatModelMessage( + role=PromptMessageRole.USER, text=QUESTION_CLASSIFIER_USER_PROMPT_2 + ) + prompt_messages.append(user_prompt_message_2) + assistant_prompt_message_2 = LLMNodeChatModelMessage( + role=PromptMessageRole.ASSISTANT, text=QUESTION_CLASSIFIER_ASSISTANT_PROMPT_2 + ) + prompt_messages.append(assistant_prompt_message_2) + user_prompt_message_3 = LLMNodeChatModelMessage( + role=PromptMessageRole.USER, + text=QUESTION_CLASSIFIER_USER_PROMPT_3.format( + input_text=input_text, + categories=json.dumps(categories, ensure_ascii=False), + classification_instructions=instruction, + ), + ) + prompt_messages.append(user_prompt_message_3) + return prompt_messages + elif model_mode == ModelMode.COMPLETION: + return LLMNodeCompletionModelPromptTemplate( + text=QUESTION_CLASSIFIER_COMPLETION_PROMPT.format( + histories=memory_str, + input_text=input_text, + categories=json.dumps(categories), + classification_instructions=instruction, + ensure_ascii=False, + ) + ) + + else: + raise InvalidModelTypeError(f"Model mode {model_mode} not support.") diff --git a/api/core/workflow/nodes/question_classifier/template_prompts.py b/api/core/workflow/nodes/question_classifier/template_prompts.py new file mode 100644 index 0000000000000000000000000000000000000000..53fc136b2c2ba56879275ea7b3adfc8bff369164 --- /dev/null +++ b/api/core/workflow/nodes/question_classifier/template_prompts.py @@ -0,0 +1,76 @@ +QUESTION_CLASSIFIER_SYSTEM_PROMPT = """ + ### Job Description', + You are a text classification engine that analyzes text data and assigns categories based on user input or automatically determined categories. + ### Task + Your task is to assign one categories ONLY to the input text and only one category may be assigned returned in the output. Additionally, you need to extract the key words from the text that are related to the classification. + ### Format + The input text is in the variable input_text. Categories are specified as a category list with two filed category_id and category_name in the variable categories. Classification instructions may be included to improve the classification accuracy. + ### Constraint + DO NOT include anything other than the JSON array in your response. + ### Memory + Here are the chat histories between human and assistant, inside XML tags. + + {histories} + +""" # noqa: E501 + +QUESTION_CLASSIFIER_USER_PROMPT_1 = """ + { "input_text": ["I recently had a great experience with your company. The service was prompt and the staff was very friendly."], + "categories": [{"category_id":"f5660049-284f-41a7-b301-fd24176a711c","category_name":"Customer Service"},{"category_id":"8d007d06-f2c9-4be5-8ff6-cd4381c13c60","category_name":"Satisfaction"},{"category_id":"5fbbbb18-9843-466d-9b8e-b9bfbb9482c8","category_name":"Sales"},{"category_id":"23623c75-7184-4a2e-8226-466c2e4631e4","category_name":"Product"}], + "classification_instructions": ["classify the text based on the feedback provided by customer"]} +""" # noqa: E501 + +QUESTION_CLASSIFIER_ASSISTANT_PROMPT_1 = """ +```json + {"keywords": ["recently", "great experience", "company", "service", "prompt", "staff", "friendly"], + "category_id": "f5660049-284f-41a7-b301-fd24176a711c", + "category_name": "Customer Service"} +``` +""" + +QUESTION_CLASSIFIER_USER_PROMPT_2 = """ + {"input_text": ["bad service, slow to bring the food"], + "categories": [{"category_id":"80fb86a0-4454-4bf5-924c-f253fdd83c02","category_name":"Food Quality"},{"category_id":"f6ff5bc3-aca0-4e4a-8627-e760d0aca78f","category_name":"Experience"},{"category_id":"cc771f63-74e7-4c61-882e-3eda9d8ba5d7","category_name":"Price"}], + "classification_instructions": []} +""" # noqa: E501 + +QUESTION_CLASSIFIER_ASSISTANT_PROMPT_2 = """ +```json + {"keywords": ["bad service", "slow", "food", "tip", "terrible", "waitresses"], + "category_id": "f6ff5bc3-aca0-4e4a-8627-e760d0aca78f", + "category_name": "Experience"} +``` +""" + +QUESTION_CLASSIFIER_USER_PROMPT_3 = """ + '{{"input_text": ["{input_text}"],', + '"categories": {categories}, ', + '"classification_instructions": ["{classification_instructions}"]}}' +""" + +QUESTION_CLASSIFIER_COMPLETION_PROMPT = """ +### Job Description +You are a text classification engine that analyzes text data and assigns categories based on user input or automatically determined categories. +### Task +Your task is to assign one categories ONLY to the input text and only one category may be assigned returned in the output. Additionally, you need to extract the key words from the text that are related to the classification. +### Format +The input text is in the variable input_text. Categories are specified as a category list with two filed category_id and category_name in the variable categories. Classification instructions may be included to improve the classification accuracy. +### Constraint +DO NOT include anything other than the JSON array in your response. +### Example +Here is the chat example between human and assistant, inside XML tags. + +User:{{"input_text": ["I recently had a great experience with your company. The service was prompt and the staff was very friendly."], "categories": [{{"category_id":"f5660049-284f-41a7-b301-fd24176a711c","category_name":"Customer Service"}},{{"category_id":"8d007d06-f2c9-4be5-8ff6-cd4381c13c60","category_name":"Satisfaction"}},{{"category_id":"5fbbbb18-9843-466d-9b8e-b9bfbb9482c8","category_name":"Sales"}},{{"category_id":"23623c75-7184-4a2e-8226-466c2e4631e4","category_name":"Product"}}], "classification_instructions": ["classify the text based on the feedback provided by customer"]}} +Assistant:{{"keywords": ["recently", "great experience", "company", "service", "prompt", "staff", "friendly"],"category_id": "f5660049-284f-41a7-b301-fd24176a711c","category_name": "Customer Service"}} +User:{{"input_text": ["bad service, slow to bring the food"], "categories": [{{"category_id":"80fb86a0-4454-4bf5-924c-f253fdd83c02","category_name":"Food Quality"}},{{"category_id":"f6ff5bc3-aca0-4e4a-8627-e760d0aca78f","category_name":"Experience"}},{{"category_id":"cc771f63-74e7-4c61-882e-3eda9d8ba5d7","category_name":"Price"}}], "classification_instructions": []}} +Assistant:{{"keywords": ["bad service", "slow", "food", "tip", "terrible", "waitresses"],"category_id": "f6ff5bc3-aca0-4e4a-8627-e760d0aca78f","category_name": "Experience"}} + +### Memory +Here are the chat histories between human and assistant, inside XML tags. + +{histories} + +### User Input +{{"input_text" : ["{input_text}"], "categories" : {categories},"classification_instruction" : ["{classification_instructions}"]}} +### Assistant Output +""" # noqa: E501 diff --git a/api/core/workflow/nodes/start/__init__.py b/api/core/workflow/nodes/start/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..54117804231aa9778c84c021802746fb39dd3c16 --- /dev/null +++ b/api/core/workflow/nodes/start/__init__.py @@ -0,0 +1,3 @@ +from .start_node import StartNode + +__all__ = ["StartNode"] diff --git a/api/core/workflow/nodes/start/entities.py b/api/core/workflow/nodes/start/entities.py new file mode 100644 index 0000000000000000000000000000000000000000..594d1b7bab8d68b340a1bfff88075e2e541f02ec --- /dev/null +++ b/api/core/workflow/nodes/start/entities.py @@ -0,0 +1,14 @@ +from collections.abc import Sequence + +from pydantic import Field + +from core.app.app_config.entities import VariableEntity +from core.workflow.nodes.base import BaseNodeData + + +class StartNodeData(BaseNodeData): + """ + Start Node Data + """ + + variables: Sequence[VariableEntity] = Field(default_factory=list) diff --git a/api/core/workflow/nodes/start/start_node.py b/api/core/workflow/nodes/start/start_node.py new file mode 100644 index 0000000000000000000000000000000000000000..a7b91e82bbdd925a3bdd0d0e44b91ed65aaf9a3e --- /dev/null +++ b/api/core/workflow/nodes/start/start_node.py @@ -0,0 +1,35 @@ +from collections.abc import Mapping, Sequence +from typing import Any + +from core.workflow.constants import SYSTEM_VARIABLE_NODE_ID +from core.workflow.entities.node_entities import NodeRunResult +from core.workflow.nodes.base import BaseNode +from core.workflow.nodes.enums import NodeType +from core.workflow.nodes.start.entities import StartNodeData +from models.workflow import WorkflowNodeExecutionStatus + + +class StartNode(BaseNode[StartNodeData]): + _node_data_cls = StartNodeData + _node_type = NodeType.START + + def _run(self) -> NodeRunResult: + node_inputs = dict(self.graph_runtime_state.variable_pool.user_inputs) + system_inputs = self.graph_runtime_state.variable_pool.system_variables + + # TODO: System variables should be directly accessible, no need for special handling + # Set system variables as node outputs. + for var in system_inputs: + node_inputs[SYSTEM_VARIABLE_NODE_ID + "." + var] = system_inputs[var] + + return NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED, inputs=node_inputs, outputs=node_inputs) + + @classmethod + def _extract_variable_selector_to_variable_mapping( + cls, + *, + graph_config: Mapping[str, Any], + node_id: str, + node_data: StartNodeData, + ) -> Mapping[str, Sequence[str]]: + return {} diff --git a/api/core/workflow/nodes/template_transform/__init__.py b/api/core/workflow/nodes/template_transform/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..43863b9d59aaf38ae8e43cff7536dce435762f98 --- /dev/null +++ b/api/core/workflow/nodes/template_transform/__init__.py @@ -0,0 +1,3 @@ +from .template_transform_node import TemplateTransformNode + +__all__ = ["TemplateTransformNode"] diff --git a/api/core/workflow/nodes/template_transform/entities.py b/api/core/workflow/nodes/template_transform/entities.py new file mode 100644 index 0000000000000000000000000000000000000000..96adff6ffaa953405b38e660c14d1813dacef238 --- /dev/null +++ b/api/core/workflow/nodes/template_transform/entities.py @@ -0,0 +1,11 @@ +from core.workflow.entities.variable_entities import VariableSelector +from core.workflow.nodes.base import BaseNodeData + + +class TemplateTransformNodeData(BaseNodeData): + """ + Code Node Data. + """ + + variables: list[VariableSelector] + template: str diff --git a/api/core/workflow/nodes/template_transform/template_transform_node.py b/api/core/workflow/nodes/template_transform/template_transform_node.py new file mode 100644 index 0000000000000000000000000000000000000000..22a1b218880db9c94547523e7a86d331e562b8b8 --- /dev/null +++ b/api/core/workflow/nodes/template_transform/template_transform_node.py @@ -0,0 +1,71 @@ +import os +from collections.abc import Mapping, Sequence +from typing import Any, Optional + +from core.helper.code_executor.code_executor import CodeExecutionError, CodeExecutor, CodeLanguage +from core.workflow.entities.node_entities import NodeRunResult +from core.workflow.nodes.base import BaseNode +from core.workflow.nodes.enums import NodeType +from core.workflow.nodes.template_transform.entities import TemplateTransformNodeData +from models.workflow import WorkflowNodeExecutionStatus + +MAX_TEMPLATE_TRANSFORM_OUTPUT_LENGTH = int(os.environ.get("TEMPLATE_TRANSFORM_MAX_LENGTH", "80000")) + + +class TemplateTransformNode(BaseNode[TemplateTransformNodeData]): + _node_data_cls = TemplateTransformNodeData + _node_type = NodeType.TEMPLATE_TRANSFORM + + @classmethod + def get_default_config(cls, filters: Optional[dict] = None) -> dict: + """ + Get default config of node. + :param filters: filter by node config parameters. + :return: + """ + return { + "type": "template-transform", + "config": {"variables": [{"variable": "arg1", "value_selector": []}], "template": "{{ arg1 }}"}, + } + + def _run(self) -> NodeRunResult: + # Get variables + variables = {} + for variable_selector in self.node_data.variables: + variable_name = variable_selector.variable + value = self.graph_runtime_state.variable_pool.get(variable_selector.value_selector) + variables[variable_name] = value.to_object() if value else None + # Run code + try: + result = CodeExecutor.execute_workflow_code_template( + language=CodeLanguage.JINJA2, code=self.node_data.template, inputs=variables + ) + except CodeExecutionError as e: + return NodeRunResult(inputs=variables, status=WorkflowNodeExecutionStatus.FAILED, error=str(e)) + + if len(result["result"]) > MAX_TEMPLATE_TRANSFORM_OUTPUT_LENGTH: + return NodeRunResult( + inputs=variables, + status=WorkflowNodeExecutionStatus.FAILED, + error=f"Output length exceeds {MAX_TEMPLATE_TRANSFORM_OUTPUT_LENGTH} characters", + ) + + return NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, inputs=variables, outputs={"output": result["result"]} + ) + + @classmethod + def _extract_variable_selector_to_variable_mapping( + cls, *, graph_config: Mapping[str, Any], node_id: str, node_data: TemplateTransformNodeData + ) -> Mapping[str, Sequence[str]]: + """ + Extract variable selector to variable mapping + :param graph_config: graph config + :param node_id: node id + :param node_data: node data + :return: + """ + return { + node_id + "." + variable_selector.variable: variable_selector.value_selector + for variable_selector in node_data.variables + } diff --git a/api/core/workflow/nodes/tool/__init__.py b/api/core/workflow/nodes/tool/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f4982e655d193fc444b9b4191cf8919f2b0a3816 --- /dev/null +++ b/api/core/workflow/nodes/tool/__init__.py @@ -0,0 +1,3 @@ +from .tool_node import ToolNode + +__all__ = ["ToolNode"] diff --git a/api/core/workflow/nodes/tool/entities.py b/api/core/workflow/nodes/tool/entities.py new file mode 100644 index 0000000000000000000000000000000000000000..9e29791481436ec0ed993a19bc8f1ae388466c77 --- /dev/null +++ b/api/core/workflow/nodes/tool/entities.py @@ -0,0 +1,54 @@ +from typing import Any, Literal, Union + +from pydantic import BaseModel, field_validator +from pydantic_core.core_schema import ValidationInfo + +from core.workflow.nodes.base import BaseNodeData + + +class ToolEntity(BaseModel): + provider_id: str + provider_type: Literal["builtin", "api", "workflow"] + provider_name: str # redundancy + tool_name: str + tool_label: str # redundancy + tool_configurations: dict[str, Any] + + @field_validator("tool_configurations", mode="before") + @classmethod + def validate_tool_configurations(cls, value, values: ValidationInfo): + if not isinstance(value, dict): + raise ValueError("tool_configurations must be a dictionary") + + for key in values.data.get("tool_configurations", {}): + value = values.data.get("tool_configurations", {}).get(key) + if not isinstance(value, str | int | float | bool): + raise ValueError(f"{key} must be a string") + + return value + + +class ToolNodeData(BaseNodeData, ToolEntity): + class ToolInput(BaseModel): + # TODO: check this type + value: Union[Any, list[str]] + type: Literal["mixed", "variable", "constant"] + + @field_validator("type", mode="before") + @classmethod + def check_type(cls, value, validation_info: ValidationInfo): + typ = value + value = validation_info.data.get("value") + if typ == "mixed" and not isinstance(value, str): + raise ValueError("value must be a string") + elif typ == "variable": + if not isinstance(value, list): + raise ValueError("value must be a list") + for val in value: + if not isinstance(val, str): + raise ValueError("value must be a list of strings") + elif typ == "constant" and not isinstance(value, str | int | float | bool): + raise ValueError("value must be a string, int, float, or bool") + return typ + + tool_parameters: dict[str, ToolInput] diff --git a/api/core/workflow/nodes/tool/exc.py b/api/core/workflow/nodes/tool/exc.py new file mode 100644 index 0000000000000000000000000000000000000000..7212e8bfc071bf922e93e4af93b581b4ebfd7288 --- /dev/null +++ b/api/core/workflow/nodes/tool/exc.py @@ -0,0 +1,16 @@ +class ToolNodeError(ValueError): + """Base exception for tool node errors.""" + + pass + + +class ToolParameterError(ToolNodeError): + """Exception raised for errors in tool parameters.""" + + pass + + +class ToolFileError(ToolNodeError): + """Exception raised for errors related to tool files.""" + + pass diff --git a/api/core/workflow/nodes/tool/tool_node.py b/api/core/workflow/nodes/tool/tool_node.py new file mode 100644 index 0000000000000000000000000000000000000000..01d07e494944b4d3496ee4c2c9066af584cfd12a --- /dev/null +++ b/api/core/workflow/nodes/tool/tool_node.py @@ -0,0 +1,308 @@ +from collections.abc import Mapping, Sequence +from typing import Any +from uuid import UUID + +from sqlalchemy import select +from sqlalchemy.orm import Session + +from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler +from core.file import File, FileTransferMethod, FileType +from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter +from core.tools.tool_engine import ToolEngine +from core.tools.utils.message_transformer import ToolFileMessageTransformer +from core.workflow.entities.node_entities import NodeRunMetadataKey, NodeRunResult +from core.workflow.entities.variable_pool import VariablePool +from core.workflow.nodes.base import BaseNode +from core.workflow.nodes.enums import NodeType +from core.workflow.utils.variable_template_parser import VariableTemplateParser +from extensions.ext_database import db +from factories import file_factory +from models import ToolFile +from models.workflow import WorkflowNodeExecutionStatus + +from .entities import ToolNodeData +from .exc import ( + ToolFileError, + ToolNodeError, + ToolParameterError, +) + + +class ToolNode(BaseNode[ToolNodeData]): + """ + Tool Node + """ + + _node_data_cls = ToolNodeData + _node_type = NodeType.TOOL + + def _run(self) -> NodeRunResult: + # fetch tool icon + tool_info = { + "provider_type": self.node_data.provider_type, + "provider_id": self.node_data.provider_id, + } + + # get tool runtime + try: + from core.tools.tool_manager import ToolManager + + tool_runtime = ToolManager.get_workflow_tool_runtime( + self.tenant_id, self.app_id, self.node_id, self.node_data, self.invoke_from + ) + except ToolNodeError as e: + return NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + inputs={}, + metadata={ + NodeRunMetadataKey.TOOL_INFO: tool_info, + }, + error=f"Failed to get tool runtime: {str(e)}", + error_type=type(e).__name__, + ) + + # get parameters + tool_parameters = tool_runtime.parameters or [] + parameters = self._generate_parameters( + tool_parameters=tool_parameters, + variable_pool=self.graph_runtime_state.variable_pool, + node_data=self.node_data, + ) + parameters_for_log = self._generate_parameters( + tool_parameters=tool_parameters, + variable_pool=self.graph_runtime_state.variable_pool, + node_data=self.node_data, + for_log=True, + ) + + try: + messages = ToolEngine.workflow_invoke( + tool=tool_runtime, + tool_parameters=parameters, + user_id=self.user_id, + workflow_tool_callback=DifyWorkflowCallbackHandler(), + workflow_call_depth=self.workflow_call_depth, + thread_pool_id=self.thread_pool_id, + ) + except ToolNodeError as e: + return NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + inputs=parameters_for_log, + metadata={ + NodeRunMetadataKey.TOOL_INFO: tool_info, + }, + error=f"Failed to invoke tool: {str(e)}", + error_type=type(e).__name__, + ) + except Exception as e: + return NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + inputs=parameters_for_log, + metadata={ + NodeRunMetadataKey.TOOL_INFO: tool_info, + }, + error=f"Failed to invoke tool: {str(e)}", + error_type="UnknownError", + ) + + # convert tool messages + plain_text, files, json = self._convert_tool_messages(messages) + + return NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + outputs={ + "text": plain_text, + "files": files, + "json": json, + }, + metadata={ + NodeRunMetadataKey.TOOL_INFO: tool_info, + }, + inputs=parameters_for_log, + ) + + def _generate_parameters( + self, + *, + tool_parameters: Sequence[ToolParameter], + variable_pool: VariablePool, + node_data: ToolNodeData, + for_log: bool = False, + ) -> Mapping[str, Any]: + """ + Generate parameters based on the given tool parameters, variable pool, and node data. + + Args: + tool_parameters (Sequence[ToolParameter]): The list of tool parameters. + variable_pool (VariablePool): The variable pool containing the variables. + node_data (ToolNodeData): The data associated with the tool node. + + Returns: + Mapping[str, Any]: A dictionary containing the generated parameters. + + """ + tool_parameters_dictionary = {parameter.name: parameter for parameter in tool_parameters} + + result: dict[str, Any] = {} + for parameter_name in node_data.tool_parameters: + parameter = tool_parameters_dictionary.get(parameter_name) + if not parameter: + result[parameter_name] = None + continue + tool_input = node_data.tool_parameters[parameter_name] + if tool_input.type == "variable": + variable = variable_pool.get(tool_input.value) + if variable is None: + raise ToolParameterError(f"Variable {tool_input.value} does not exist") + parameter_value = variable.value + elif tool_input.type in {"mixed", "constant"}: + segment_group = variable_pool.convert_template(str(tool_input.value)) + parameter_value = segment_group.log if for_log else segment_group.text + else: + raise ToolParameterError(f"Unknown tool input type '{tool_input.type}'") + result[parameter_name] = parameter_value + + return result + + def _convert_tool_messages( + self, + messages: list[ToolInvokeMessage], + ): + """ + Convert ToolInvokeMessages into tuple[plain_text, files] + """ + # transform message and handle file storage + messages = ToolFileMessageTransformer.transform_tool_invoke_messages( + messages=messages, + user_id=self.user_id, + tenant_id=self.tenant_id, + conversation_id=None, + ) + # extract plain text and files + files = self._extract_tool_response_binary(messages) + plain_text = self._extract_tool_response_text(messages) + json = self._extract_tool_response_json(messages) + + return plain_text, files, json + + def _extract_tool_response_binary(self, tool_response: list[ToolInvokeMessage]) -> list[File]: + """ + Extract tool response binary + """ + result = [] + for response in tool_response: + if response.type in {ToolInvokeMessage.MessageType.IMAGE_LINK, ToolInvokeMessage.MessageType.IMAGE}: + url = str(response.message) if response.message else None + tool_file_id = str(url).split("/")[-1].split(".")[0] + transfer_method = response.meta.get("transfer_method", FileTransferMethod.TOOL_FILE) + + with Session(db.engine) as session: + stmt = select(ToolFile).where(ToolFile.id == tool_file_id) + tool_file = session.scalar(stmt) + if tool_file is None: + raise ToolFileError(f"Tool file {tool_file_id} does not exist") + + mapping = { + "tool_file_id": tool_file_id, + "type": FileType.IMAGE, + "transfer_method": transfer_method, + "url": url, + } + file = file_factory.build_from_mapping( + mapping=mapping, + tenant_id=self.tenant_id, + ) + result.append(file) + elif response.type == ToolInvokeMessage.MessageType.BLOB: + tool_file_id = str(response.message).split("/")[-1].split(".")[0] + with Session(db.engine) as session: + stmt = select(ToolFile).where(ToolFile.id == tool_file_id) + tool_file = session.scalar(stmt) + if tool_file is None: + raise ValueError(f"tool file {tool_file_id} not exists") + mapping = { + "tool_file_id": tool_file_id, + "transfer_method": FileTransferMethod.TOOL_FILE, + } + file = file_factory.build_from_mapping( + mapping=mapping, + tenant_id=self.tenant_id, + ) + result.append(file) + elif response.type == ToolInvokeMessage.MessageType.LINK: + url = str(response.message) + transfer_method = FileTransferMethod.TOOL_FILE + tool_file_id = url.split("/")[-1].split(".")[0] + try: + UUID(tool_file_id) + except ValueError: + raise ToolFileError(f"cannot extract tool file id from url {url}") + with Session(db.engine) as session: + stmt = select(ToolFile).where(ToolFile.id == tool_file_id) + tool_file = session.scalar(stmt) + if tool_file is None: + raise ToolFileError(f"Tool file {tool_file_id} does not exist") + mapping = { + "tool_file_id": tool_file_id, + "transfer_method": transfer_method, + "url": url, + } + file = file_factory.build_from_mapping( + mapping=mapping, + tenant_id=self.tenant_id, + ) + result.append(file) + + elif response.type == ToolInvokeMessage.MessageType.FILE: + assert response.meta is not None + result.append(response.meta["file"]) + + return result + + def _extract_tool_response_text(self, tool_response: list[ToolInvokeMessage]) -> str: + """ + Extract tool response text + """ + return "\n".join( + [ + str(message.message) + if message.type == ToolInvokeMessage.MessageType.TEXT + else f"Link: {str(message.message)}" + for message in tool_response + if message.type in {ToolInvokeMessage.MessageType.TEXT, ToolInvokeMessage.MessageType.LINK} + ] + ) + + def _extract_tool_response_json(self, tool_response: list[ToolInvokeMessage]): + return [message.message for message in tool_response if message.type == ToolInvokeMessage.MessageType.JSON] + + @classmethod + def _extract_variable_selector_to_variable_mapping( + cls, + *, + graph_config: Mapping[str, Any], + node_id: str, + node_data: ToolNodeData, + ) -> Mapping[str, Sequence[str]]: + """ + Extract variable selector to variable mapping + :param graph_config: graph config + :param node_id: node id + :param node_data: node data + :return: + """ + result = {} + for parameter_name in node_data.tool_parameters: + input = node_data.tool_parameters[parameter_name] + if input.type == "mixed": + selectors = VariableTemplateParser(str(input.value)).extract_variable_selectors() + for selector in selectors: + result[selector.variable] = selector.value_selector + elif input.type == "variable": + result[parameter_name] = input.value + elif input.type == "constant": + pass + + result = {node_id + "." + key: value for key, value in result.items()} + + return result diff --git a/api/core/workflow/nodes/variable_aggregator/__init__.py b/api/core/workflow/nodes/variable_aggregator/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0b6bf2a5b62ada0f5f493939dc3f7e88426ea20c --- /dev/null +++ b/api/core/workflow/nodes/variable_aggregator/__init__.py @@ -0,0 +1,3 @@ +from .variable_aggregator_node import VariableAggregatorNode + +__all__ = ["VariableAggregatorNode"] diff --git a/api/core/workflow/nodes/variable_aggregator/entities.py b/api/core/workflow/nodes/variable_aggregator/entities.py new file mode 100644 index 0000000000000000000000000000000000000000..71a930e6b0a5cb1a0c1cc508bedacc49cd4c7d37 --- /dev/null +++ b/api/core/workflow/nodes/variable_aggregator/entities.py @@ -0,0 +1,35 @@ +from typing import Literal, Optional + +from pydantic import BaseModel + +from core.workflow.nodes.base import BaseNodeData + + +class AdvancedSettings(BaseModel): + """ + Advanced setting. + """ + + group_enabled: bool + + class Group(BaseModel): + """ + Group. + """ + + output_type: Literal["string", "number", "object", "array[string]", "array[number]", "array[object]"] + variables: list[list[str]] + group_name: str + + groups: list[Group] + + +class VariableAssignerNodeData(BaseNodeData): + """ + Knowledge retrieval Node Data. + """ + + type: str = "variable-assigner" + output_type: str + variables: list[list[str]] + advanced_settings: Optional[AdvancedSettings] = None diff --git a/api/core/workflow/nodes/variable_aggregator/variable_aggregator_node.py b/api/core/workflow/nodes/variable_aggregator/variable_aggregator_node.py new file mode 100644 index 0000000000000000000000000000000000000000..031a7b83095541d79e9b9fb8c6dea6154fbe7854 --- /dev/null +++ b/api/core/workflow/nodes/variable_aggregator/variable_aggregator_node.py @@ -0,0 +1,51 @@ +from collections.abc import Mapping, Sequence +from typing import Any + +from core.workflow.entities.node_entities import NodeRunResult +from core.workflow.nodes.base import BaseNode +from core.workflow.nodes.enums import NodeType +from core.workflow.nodes.variable_aggregator.entities import VariableAssignerNodeData +from models.workflow import WorkflowNodeExecutionStatus + + +class VariableAggregatorNode(BaseNode[VariableAssignerNodeData]): + _node_data_cls = VariableAssignerNodeData + _node_type = NodeType.VARIABLE_AGGREGATOR + + def _run(self) -> NodeRunResult: + # Get variables + outputs = {} + inputs = {} + + if not self.node_data.advanced_settings or not self.node_data.advanced_settings.group_enabled: + for selector in self.node_data.variables: + variable = self.graph_runtime_state.variable_pool.get(selector) + if variable is not None: + outputs = {"output": variable.to_object()} + + inputs = {".".join(selector[1:]): variable.to_object()} + break + else: + for group in self.node_data.advanced_settings.groups: + for selector in group.variables: + variable = self.graph_runtime_state.variable_pool.get(selector) + + if variable is not None: + outputs[group.group_name] = {"output": variable.to_object()} + inputs[".".join(selector[1:])] = variable.to_object() + break + + return NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED, outputs=outputs, inputs=inputs) + + @classmethod + def _extract_variable_selector_to_variable_mapping( + cls, *, graph_config: Mapping[str, Any], node_id: str, node_data: VariableAssignerNodeData + ) -> Mapping[str, Sequence[str]]: + """ + Extract variable selector to variable mapping + :param graph_config: graph config + :param node_id: node id + :param node_data: node data + :return: + """ + return {} diff --git a/api/core/workflow/nodes/variable_assigner/__init__.py b/api/core/workflow/nodes/variable_assigner/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/workflow/nodes/variable_assigner/common/__init__.py b/api/core/workflow/nodes/variable_assigner/common/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/workflow/nodes/variable_assigner/common/exc.py b/api/core/workflow/nodes/variable_assigner/common/exc.py new file mode 100644 index 0000000000000000000000000000000000000000..f8dbedc2901c9f4c4b79d5d6093011c6893447a8 --- /dev/null +++ b/api/core/workflow/nodes/variable_assigner/common/exc.py @@ -0,0 +1,4 @@ +class VariableOperatorNodeError(ValueError): + """Base error type, don't use directly.""" + + pass diff --git a/api/core/workflow/nodes/variable_assigner/common/helpers.py b/api/core/workflow/nodes/variable_assigner/common/helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..8031b57fa82892b3dd4002a2527da388bdedd0a5 --- /dev/null +++ b/api/core/workflow/nodes/variable_assigner/common/helpers.py @@ -0,0 +1,19 @@ +from sqlalchemy import select +from sqlalchemy.orm import Session + +from core.variables import Variable +from core.workflow.nodes.variable_assigner.common.exc import VariableOperatorNodeError +from extensions.ext_database import db +from models import ConversationVariable + + +def update_conversation_variable(conversation_id: str, variable: Variable): + stmt = select(ConversationVariable).where( + ConversationVariable.id == variable.id, ConversationVariable.conversation_id == conversation_id + ) + with Session(db.engine) as session: + row = session.scalar(stmt) + if not row: + raise VariableOperatorNodeError("conversation variable not found in the database") + row.data = variable.model_dump_json() + session.commit() diff --git a/api/core/workflow/nodes/variable_assigner/v1/__init__.py b/api/core/workflow/nodes/variable_assigner/v1/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7eb1428e50370cba578acd12b32aba4dfba07a06 --- /dev/null +++ b/api/core/workflow/nodes/variable_assigner/v1/__init__.py @@ -0,0 +1,3 @@ +from .node import VariableAssignerNode + +__all__ = ["VariableAssignerNode"] diff --git a/api/core/workflow/nodes/variable_assigner/v1/node.py b/api/core/workflow/nodes/variable_assigner/v1/node.py new file mode 100644 index 0000000000000000000000000000000000000000..9acc76f326eec9a5a7f560f3c53f56fb07e6d91d --- /dev/null +++ b/api/core/workflow/nodes/variable_assigner/v1/node.py @@ -0,0 +1,75 @@ +from core.variables import SegmentType, Variable +from core.workflow.entities.node_entities import NodeRunResult +from core.workflow.nodes.base import BaseNode, BaseNodeData +from core.workflow.nodes.enums import NodeType +from core.workflow.nodes.variable_assigner.common import helpers as common_helpers +from core.workflow.nodes.variable_assigner.common.exc import VariableOperatorNodeError +from factories import variable_factory +from models.workflow import WorkflowNodeExecutionStatus + +from .node_data import VariableAssignerData, WriteMode + + +class VariableAssignerNode(BaseNode[VariableAssignerData]): + _node_data_cls: type[BaseNodeData] = VariableAssignerData + _node_type = NodeType.VARIABLE_ASSIGNER + + def _run(self) -> NodeRunResult: + # Should be String, Number, Object, ArrayString, ArrayNumber, ArrayObject + original_variable = self.graph_runtime_state.variable_pool.get(self.node_data.assigned_variable_selector) + if not isinstance(original_variable, Variable): + raise VariableOperatorNodeError("assigned variable not found") + + match self.node_data.write_mode: + case WriteMode.OVER_WRITE: + income_value = self.graph_runtime_state.variable_pool.get(self.node_data.input_variable_selector) + if not income_value: + raise VariableOperatorNodeError("input value not found") + updated_variable = original_variable.model_copy(update={"value": income_value.value}) + + case WriteMode.APPEND: + income_value = self.graph_runtime_state.variable_pool.get(self.node_data.input_variable_selector) + if not income_value: + raise VariableOperatorNodeError("input value not found") + updated_value = original_variable.value + [income_value.value] + updated_variable = original_variable.model_copy(update={"value": updated_value}) + + case WriteMode.CLEAR: + income_value = get_zero_value(original_variable.value_type) + if income_value is None: + raise VariableOperatorNodeError("income value not found") + updated_variable = original_variable.model_copy(update={"value": income_value.to_object()}) + + case _: + raise VariableOperatorNodeError(f"unsupported write mode: {self.node_data.write_mode}") + + # Over write the variable. + self.graph_runtime_state.variable_pool.add(self.node_data.assigned_variable_selector, updated_variable) + + # TODO: Move database operation to the pipeline. + # Update conversation variable. + conversation_id = self.graph_runtime_state.variable_pool.get(["sys", "conversation_id"]) + if not conversation_id: + raise VariableOperatorNodeError("conversation_id not found") + common_helpers.update_conversation_variable(conversation_id=conversation_id.text, variable=updated_variable) + + return NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + inputs={ + "value": income_value.to_object(), + }, + ) + + +def get_zero_value(t: SegmentType): + match t: + case SegmentType.ARRAY_OBJECT | SegmentType.ARRAY_STRING | SegmentType.ARRAY_NUMBER: + return variable_factory.build_segment([]) + case SegmentType.OBJECT: + return variable_factory.build_segment({}) + case SegmentType.STRING: + return variable_factory.build_segment("") + case SegmentType.NUMBER: + return variable_factory.build_segment(0) + case _: + raise VariableOperatorNodeError(f"unsupported variable type: {t}") diff --git a/api/core/workflow/nodes/variable_assigner/v1/node_data.py b/api/core/workflow/nodes/variable_assigner/v1/node_data.py new file mode 100644 index 0000000000000000000000000000000000000000..9734d64712067c03aee7d386843990bb9d9cc189 --- /dev/null +++ b/api/core/workflow/nodes/variable_assigner/v1/node_data.py @@ -0,0 +1,16 @@ +from collections.abc import Sequence +from enum import StrEnum + +from core.workflow.nodes.base import BaseNodeData + + +class WriteMode(StrEnum): + OVER_WRITE = "over-write" + APPEND = "append" + CLEAR = "clear" + + +class VariableAssignerData(BaseNodeData): + assigned_variable_selector: Sequence[str] + write_mode: WriteMode + input_variable_selector: Sequence[str] diff --git a/api/core/workflow/nodes/variable_assigner/v2/__init__.py b/api/core/workflow/nodes/variable_assigner/v2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7eb1428e50370cba578acd12b32aba4dfba07a06 --- /dev/null +++ b/api/core/workflow/nodes/variable_assigner/v2/__init__.py @@ -0,0 +1,3 @@ +from .node import VariableAssignerNode + +__all__ = ["VariableAssignerNode"] diff --git a/api/core/workflow/nodes/variable_assigner/v2/constants.py b/api/core/workflow/nodes/variable_assigner/v2/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..3797bfa77a1d6258bbcc0507d3790d72ec9ae3fb --- /dev/null +++ b/api/core/workflow/nodes/variable_assigner/v2/constants.py @@ -0,0 +1,11 @@ +from core.variables import SegmentType + +EMPTY_VALUE_MAPPING = { + SegmentType.STRING: "", + SegmentType.NUMBER: 0, + SegmentType.OBJECT: {}, + SegmentType.ARRAY_ANY: [], + SegmentType.ARRAY_STRING: [], + SegmentType.ARRAY_NUMBER: [], + SegmentType.ARRAY_OBJECT: [], +} diff --git a/api/core/workflow/nodes/variable_assigner/v2/entities.py b/api/core/workflow/nodes/variable_assigner/v2/entities.py new file mode 100644 index 0000000000000000000000000000000000000000..01df33b6d448fc648b35ba403d859d44e146fe72 --- /dev/null +++ b/api/core/workflow/nodes/variable_assigner/v2/entities.py @@ -0,0 +1,20 @@ +from collections.abc import Sequence +from typing import Any + +from pydantic import BaseModel + +from core.workflow.nodes.base import BaseNodeData + +from .enums import InputType, Operation + + +class VariableOperationItem(BaseModel): + variable_selector: Sequence[str] + input_type: InputType + operation: Operation + value: Any | None = None + + +class VariableAssignerNodeData(BaseNodeData): + version: str = "2" + items: Sequence[VariableOperationItem] diff --git a/api/core/workflow/nodes/variable_assigner/v2/enums.py b/api/core/workflow/nodes/variable_assigner/v2/enums.py new file mode 100644 index 0000000000000000000000000000000000000000..36cf68aa1913fd3deb060644d103aa3885d72b33 --- /dev/null +++ b/api/core/workflow/nodes/variable_assigner/v2/enums.py @@ -0,0 +1,18 @@ +from enum import StrEnum + + +class Operation(StrEnum): + OVER_WRITE = "over-write" + CLEAR = "clear" + APPEND = "append" + EXTEND = "extend" + SET = "set" + ADD = "+=" + SUBTRACT = "-=" + MULTIPLY = "*=" + DIVIDE = "/=" + + +class InputType(StrEnum): + VARIABLE = "variable" + CONSTANT = "constant" diff --git a/api/core/workflow/nodes/variable_assigner/v2/exc.py b/api/core/workflow/nodes/variable_assigner/v2/exc.py new file mode 100644 index 0000000000000000000000000000000000000000..b67af6d73c44d5f49951f7e42ea04cbe6f6f81ce --- /dev/null +++ b/api/core/workflow/nodes/variable_assigner/v2/exc.py @@ -0,0 +1,31 @@ +from collections.abc import Sequence +from typing import Any + +from core.workflow.nodes.variable_assigner.common.exc import VariableOperatorNodeError + +from .enums import InputType, Operation + + +class OperationNotSupportedError(VariableOperatorNodeError): + def __init__(self, *, operation: Operation, variable_type: str): + super().__init__(f"Operation {operation} is not supported for type {variable_type}") + + +class InputTypeNotSupportedError(VariableOperatorNodeError): + def __init__(self, *, input_type: InputType, operation: Operation): + super().__init__(f"Input type {input_type} is not supported for operation {operation}") + + +class VariableNotFoundError(VariableOperatorNodeError): + def __init__(self, *, variable_selector: Sequence[str]): + super().__init__(f"Variable {variable_selector} not found") + + +class InvalidInputValueError(VariableOperatorNodeError): + def __init__(self, *, value: Any): + super().__init__(f"Invalid input value {value}") + + +class ConversationIDNotFoundError(VariableOperatorNodeError): + def __init__(self): + super().__init__("conversation_id not found") diff --git a/api/core/workflow/nodes/variable_assigner/v2/helpers.py b/api/core/workflow/nodes/variable_assigner/v2/helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..a86c7eb94a50afda2ce6f2b2128749668af264d0 --- /dev/null +++ b/api/core/workflow/nodes/variable_assigner/v2/helpers.py @@ -0,0 +1,91 @@ +from typing import Any + +from core.variables import SegmentType + +from .enums import Operation + + +def is_operation_supported(*, variable_type: SegmentType, operation: Operation): + match operation: + case Operation.OVER_WRITE | Operation.CLEAR: + return True + case Operation.SET: + return variable_type in {SegmentType.OBJECT, SegmentType.STRING, SegmentType.NUMBER} + case Operation.ADD | Operation.SUBTRACT | Operation.MULTIPLY | Operation.DIVIDE: + # Only number variable can be added, subtracted, multiplied or divided + return variable_type == SegmentType.NUMBER + case Operation.APPEND | Operation.EXTEND: + # Only array variable can be appended or extended + return variable_type in { + SegmentType.ARRAY_ANY, + SegmentType.ARRAY_OBJECT, + SegmentType.ARRAY_STRING, + SegmentType.ARRAY_NUMBER, + SegmentType.ARRAY_FILE, + } + case _: + return False + + +def is_variable_input_supported(*, operation: Operation): + if operation in {Operation.SET, Operation.ADD, Operation.SUBTRACT, Operation.MULTIPLY, Operation.DIVIDE}: + return False + return True + + +def is_constant_input_supported(*, variable_type: SegmentType, operation: Operation): + match variable_type: + case SegmentType.STRING | SegmentType.OBJECT: + return operation in {Operation.OVER_WRITE, Operation.SET} + case SegmentType.NUMBER: + return operation in { + Operation.OVER_WRITE, + Operation.SET, + Operation.ADD, + Operation.SUBTRACT, + Operation.MULTIPLY, + Operation.DIVIDE, + } + case _: + return False + + +def is_input_value_valid(*, variable_type: SegmentType, operation: Operation, value: Any): + if operation == Operation.CLEAR: + return True + match variable_type: + case SegmentType.STRING: + return isinstance(value, str) + + case SegmentType.NUMBER: + if not isinstance(value, int | float): + return False + if operation == Operation.DIVIDE and value == 0: + return False + return True + + case SegmentType.OBJECT: + return isinstance(value, dict) + + # Array & Append + case SegmentType.ARRAY_ANY if operation == Operation.APPEND: + return isinstance(value, str | float | int | dict) + case SegmentType.ARRAY_STRING if operation == Operation.APPEND: + return isinstance(value, str) + case SegmentType.ARRAY_NUMBER if operation == Operation.APPEND: + return isinstance(value, int | float) + case SegmentType.ARRAY_OBJECT if operation == Operation.APPEND: + return isinstance(value, dict) + + # Array & Extend / Overwrite + case SegmentType.ARRAY_ANY if operation in {Operation.EXTEND, Operation.OVER_WRITE}: + return isinstance(value, list) and all(isinstance(item, str | float | int | dict) for item in value) + case SegmentType.ARRAY_STRING if operation in {Operation.EXTEND, Operation.OVER_WRITE}: + return isinstance(value, list) and all(isinstance(item, str) for item in value) + case SegmentType.ARRAY_NUMBER if operation in {Operation.EXTEND, Operation.OVER_WRITE}: + return isinstance(value, list) and all(isinstance(item, int | float) for item in value) + case SegmentType.ARRAY_OBJECT if operation in {Operation.EXTEND, Operation.OVER_WRITE}: + return isinstance(value, list) and all(isinstance(item, dict) for item in value) + + case _: + return False diff --git a/api/core/workflow/nodes/variable_assigner/v2/node.py b/api/core/workflow/nodes/variable_assigner/v2/node.py new file mode 100644 index 0000000000000000000000000000000000000000..afa5656f46e692026ebea8d1bbc5d856546ffa4e --- /dev/null +++ b/api/core/workflow/nodes/variable_assigner/v2/node.py @@ -0,0 +1,167 @@ +import json +from collections.abc import Sequence +from typing import Any, cast + +from core.variables import SegmentType, Variable +from core.workflow.constants import CONVERSATION_VARIABLE_NODE_ID +from core.workflow.entities.node_entities import NodeRunResult +from core.workflow.nodes.base import BaseNode +from core.workflow.nodes.enums import NodeType +from core.workflow.nodes.variable_assigner.common import helpers as common_helpers +from core.workflow.nodes.variable_assigner.common.exc import VariableOperatorNodeError +from models.workflow import WorkflowNodeExecutionStatus + +from . import helpers +from .constants import EMPTY_VALUE_MAPPING +from .entities import VariableAssignerNodeData +from .enums import InputType, Operation +from .exc import ( + ConversationIDNotFoundError, + InputTypeNotSupportedError, + InvalidInputValueError, + OperationNotSupportedError, + VariableNotFoundError, +) + + +class VariableAssignerNode(BaseNode[VariableAssignerNodeData]): + _node_data_cls = VariableAssignerNodeData + _node_type = NodeType.VARIABLE_ASSIGNER + + def _run(self) -> NodeRunResult: + inputs = self.node_data.model_dump() + process_data: dict[str, Any] = {} + # NOTE: This node has no outputs + updated_variable_selectors: list[Sequence[str]] = [] + + try: + for item in self.node_data.items: + variable = self.graph_runtime_state.variable_pool.get(item.variable_selector) + + # ==================== Validation Part + + # Check if variable exists + if not isinstance(variable, Variable): + raise VariableNotFoundError(variable_selector=item.variable_selector) + + # Check if operation is supported + if not helpers.is_operation_supported(variable_type=variable.value_type, operation=item.operation): + raise OperationNotSupportedError(operation=item.operation, variable_type=variable.value_type) + + # Check if variable input is supported + if item.input_type == InputType.VARIABLE and not helpers.is_variable_input_supported( + operation=item.operation + ): + raise InputTypeNotSupportedError(input_type=InputType.VARIABLE, operation=item.operation) + + # Check if constant input is supported + if item.input_type == InputType.CONSTANT and not helpers.is_constant_input_supported( + variable_type=variable.value_type, operation=item.operation + ): + raise InputTypeNotSupportedError(input_type=InputType.CONSTANT, operation=item.operation) + + # Get value from variable pool + if ( + item.input_type == InputType.VARIABLE + and item.operation != Operation.CLEAR + and item.value is not None + ): + value = self.graph_runtime_state.variable_pool.get(item.value) + if value is None: + raise VariableNotFoundError(variable_selector=item.value) + # Skip if value is NoneSegment + if value.value_type == SegmentType.NONE: + continue + item.value = value.value + + # If set string / bytes / bytearray to object, try convert string to object. + if ( + item.operation == Operation.SET + and variable.value_type == SegmentType.OBJECT + and isinstance(item.value, str | bytes | bytearray) + ): + try: + item.value = json.loads(item.value) + except json.JSONDecodeError: + raise InvalidInputValueError(value=item.value) + + # Check if input value is valid + if not helpers.is_input_value_valid( + variable_type=variable.value_type, operation=item.operation, value=item.value + ): + raise InvalidInputValueError(value=item.value) + + # ==================== Execution Part + + updated_value = self._handle_item( + variable=variable, + operation=item.operation, + value=item.value, + ) + variable = variable.model_copy(update={"value": updated_value}) + self.graph_runtime_state.variable_pool.add(variable.selector, variable) + updated_variable_selectors.append(variable.selector) + except VariableOperatorNodeError as e: + return NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + inputs=inputs, + process_data=process_data, + error=str(e), + ) + + # The `updated_variable_selectors` is a list contains list[str] which not hashable, + # remove the duplicated items first. + updated_variable_selectors = list(set(map(tuple, updated_variable_selectors))) + + # Update variables + for selector in updated_variable_selectors: + variable = self.graph_runtime_state.variable_pool.get(selector) + if not isinstance(variable, Variable): + raise VariableNotFoundError(variable_selector=selector) + process_data[variable.name] = variable.value + + if variable.selector[0] == CONVERSATION_VARIABLE_NODE_ID: + conversation_id = self.graph_runtime_state.variable_pool.get(["sys", "conversation_id"]) + if not conversation_id: + raise ConversationIDNotFoundError + else: + conversation_id = conversation_id.value + common_helpers.update_conversation_variable( + conversation_id=cast(str, conversation_id), + variable=variable, + ) + + return NodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + inputs=inputs, + process_data=process_data, + ) + + def _handle_item( + self, + *, + variable: Variable, + operation: Operation, + value: Any, + ): + match operation: + case Operation.OVER_WRITE: + return value + case Operation.CLEAR: + return EMPTY_VALUE_MAPPING[variable.value_type] + case Operation.APPEND: + return variable.value + [value] + case Operation.EXTEND: + return variable.value + value + case Operation.SET: + return value + case Operation.ADD: + return variable.value + value + case Operation.SUBTRACT: + return variable.value - value + case Operation.MULTIPLY: + return variable.value * value + case Operation.DIVIDE: + return variable.value / value + case _: + raise OperationNotSupportedError(operation=operation, variable_type=variable.value_type) diff --git a/api/core/workflow/utils/__init__.py b/api/core/workflow/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/workflow/utils/condition/__init__.py b/api/core/workflow/utils/condition/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/workflow/utils/condition/entities.py b/api/core/workflow/utils/condition/entities.py new file mode 100644 index 0000000000000000000000000000000000000000..799c735f5409ee02367eb48739575cf11d6c3737 --- /dev/null +++ b/api/core/workflow/utils/condition/entities.py @@ -0,0 +1,49 @@ +from collections.abc import Sequence +from typing import Literal + +from pydantic import BaseModel, Field + +SupportedComparisonOperator = Literal[ + # for string or array + "contains", + "not contains", + "start with", + "end with", + "is", + "is not", + "empty", + "not empty", + "in", + "not in", + "all of", + # for number + "=", + "≠", + ">", + "<", + "≥", + "≤", + "null", + "not null", + # for file + "exists", + "not exists", +] + + +class SubCondition(BaseModel): + key: str + comparison_operator: SupportedComparisonOperator + value: str | Sequence[str] | None = None + + +class SubVariableCondition(BaseModel): + logical_operator: Literal["and", "or"] + conditions: list[SubCondition] = Field(default=list) + + +class Condition(BaseModel): + variable_selector: list[str] + comparison_operator: SupportedComparisonOperator + value: str | Sequence[str] | None = None + sub_variable_condition: SubVariableCondition | None = None diff --git a/api/core/workflow/utils/condition/processor.py b/api/core/workflow/utils/condition/processor.py new file mode 100644 index 0000000000000000000000000000000000000000..19473f39d2299af6599d7adefd46640d72722024 --- /dev/null +++ b/api/core/workflow/utils/condition/processor.py @@ -0,0 +1,385 @@ +from collections.abc import Sequence +from typing import Any, Literal + +from core.file import FileAttribute, file_manager +from core.variables import ArrayFileSegment +from core.workflow.entities.variable_pool import VariablePool + +from .entities import Condition, SubCondition, SupportedComparisonOperator + + +class ConditionProcessor: + def process_conditions( + self, + *, + variable_pool: VariablePool, + conditions: Sequence[Condition], + operator: Literal["and", "or"], + ): + input_conditions = [] + group_results = [] + + for condition in conditions: + variable = variable_pool.get(condition.variable_selector) + if variable is None: + raise ValueError(f"Variable {condition.variable_selector} not found") + + if isinstance(variable, ArrayFileSegment) and condition.comparison_operator in { + "contains", + "not contains", + "all of", + }: + # check sub conditions + if not condition.sub_variable_condition: + raise ValueError("Sub variable is required") + result = _process_sub_conditions( + variable=variable, + sub_conditions=condition.sub_variable_condition.conditions, + operator=condition.sub_variable_condition.logical_operator, + ) + elif condition.comparison_operator in { + "exists", + "not exists", + }: + result = _evaluate_condition( + value=variable.value, + operator=condition.comparison_operator, + expected=None, + ) + else: + actual_value = variable.value if variable else None + expected_value = condition.value + if isinstance(expected_value, str): + expected_value = variable_pool.convert_template(expected_value).text + input_conditions.append( + { + "actual_value": actual_value, + "expected_value": expected_value, + "comparison_operator": condition.comparison_operator, + } + ) + result = _evaluate_condition( + value=actual_value, + operator=condition.comparison_operator, + expected=expected_value, + ) + group_results.append(result) + + final_result = all(group_results) if operator == "and" else any(group_results) + return input_conditions, group_results, final_result + + +def _evaluate_condition( + *, + operator: SupportedComparisonOperator, + value: Any, + expected: str | Sequence[str] | None, +) -> bool: + match operator: + case "contains": + return _assert_contains(value=value, expected=expected) + case "not contains": + return _assert_not_contains(value=value, expected=expected) + case "start with": + return _assert_start_with(value=value, expected=expected) + case "end with": + return _assert_end_with(value=value, expected=expected) + case "is": + return _assert_is(value=value, expected=expected) + case "is not": + return _assert_is_not(value=value, expected=expected) + case "empty": + return _assert_empty(value=value) + case "not empty": + return _assert_not_empty(value=value) + case "=": + return _assert_equal(value=value, expected=expected) + case "≠": + return _assert_not_equal(value=value, expected=expected) + case ">": + return _assert_greater_than(value=value, expected=expected) + case "<": + return _assert_less_than(value=value, expected=expected) + case "≥": + return _assert_greater_than_or_equal(value=value, expected=expected) + case "≤": + return _assert_less_than_or_equal(value=value, expected=expected) + case "null": + return _assert_null(value=value) + case "not null": + return _assert_not_null(value=value) + case "in": + return _assert_in(value=value, expected=expected) + case "not in": + return _assert_not_in(value=value, expected=expected) + case "all of" if isinstance(expected, list): + return _assert_all_of(value=value, expected=expected) + case "exists": + return _assert_exists(value=value) + case "not exists": + return _assert_not_exists(value=value) + case _: + raise ValueError(f"Unsupported operator: {operator}") + + +def _assert_contains(*, value: Any, expected: Any) -> bool: + if not value: + return False + + if not isinstance(value, str | list): + raise ValueError("Invalid actual value type: string or array") + + if expected not in value: + return False + return True + + +def _assert_not_contains(*, value: Any, expected: Any) -> bool: + if not value: + return True + + if not isinstance(value, str | list): + raise ValueError("Invalid actual value type: string or array") + + if expected in value: + return False + return True + + +def _assert_start_with(*, value: Any, expected: Any) -> bool: + if not value: + return False + + if not isinstance(value, str): + raise ValueError("Invalid actual value type: string") + + if not value.startswith(expected): + return False + return True + + +def _assert_end_with(*, value: Any, expected: Any) -> bool: + if not value: + return False + + if not isinstance(value, str): + raise ValueError("Invalid actual value type: string") + + if not value.endswith(expected): + return False + return True + + +def _assert_is(*, value: Any, expected: Any) -> bool: + if value is None: + return False + + if not isinstance(value, str): + raise ValueError("Invalid actual value type: string") + + if value != expected: + return False + return True + + +def _assert_is_not(*, value: Any, expected: Any) -> bool: + if value is None: + return False + + if not isinstance(value, str): + raise ValueError("Invalid actual value type: string") + + if value == expected: + return False + return True + + +def _assert_empty(*, value: Any) -> bool: + if not value: + return True + return False + + +def _assert_not_empty(*, value: Any) -> bool: + if value: + return True + return False + + +def _assert_equal(*, value: Any, expected: Any) -> bool: + if value is None: + return False + + if not isinstance(value, int | float): + raise ValueError("Invalid actual value type: number") + + if isinstance(value, int): + expected = int(expected) + else: + expected = float(expected) + + if value != expected: + return False + return True + + +def _assert_not_equal(*, value: Any, expected: Any) -> bool: + if value is None: + return False + + if not isinstance(value, int | float): + raise ValueError("Invalid actual value type: number") + + if isinstance(value, int): + expected = int(expected) + else: + expected = float(expected) + + if value == expected: + return False + return True + + +def _assert_greater_than(*, value: Any, expected: Any) -> bool: + if value is None: + return False + + if not isinstance(value, int | float): + raise ValueError("Invalid actual value type: number") + + if isinstance(value, int): + expected = int(expected) + else: + expected = float(expected) + + if value <= expected: + return False + return True + + +def _assert_less_than(*, value: Any, expected: Any) -> bool: + if value is None: + return False + + if not isinstance(value, int | float): + raise ValueError("Invalid actual value type: number") + + if isinstance(value, int): + expected = int(expected) + else: + expected = float(expected) + + if value >= expected: + return False + return True + + +def _assert_greater_than_or_equal(*, value: Any, expected: Any) -> bool: + if value is None: + return False + + if not isinstance(value, int | float): + raise ValueError("Invalid actual value type: number") + + if isinstance(value, int): + expected = int(expected) + else: + expected = float(expected) + + if value < expected: + return False + return True + + +def _assert_less_than_or_equal(*, value: Any, expected: Any) -> bool: + if value is None: + return False + + if not isinstance(value, int | float): + raise ValueError("Invalid actual value type: number") + + if isinstance(value, int): + expected = int(expected) + else: + expected = float(expected) + + if value > expected: + return False + return True + + +def _assert_null(*, value: Any) -> bool: + if value is None: + return True + return False + + +def _assert_not_null(*, value: Any) -> bool: + if value is not None: + return True + return False + + +def _assert_in(*, value: Any, expected: Any) -> bool: + if not value: + return False + + if not isinstance(expected, list): + raise ValueError("Invalid expected value type: array") + + if value not in expected: + return False + return True + + +def _assert_not_in(*, value: Any, expected: Any) -> bool: + if not value: + return True + + if not isinstance(expected, list): + raise ValueError("Invalid expected value type: array") + + if value in expected: + return False + return True + + +def _assert_all_of(*, value: Any, expected: Sequence[str]) -> bool: + if not value: + return False + + if not all(item in value for item in expected): + return False + return True + + +def _assert_exists(*, value: Any) -> bool: + return value is not None + + +def _assert_not_exists(*, value: Any) -> bool: + return value is None + + +def _process_sub_conditions( + variable: ArrayFileSegment, + sub_conditions: Sequence[SubCondition], + operator: Literal["and", "or"], +) -> bool: + files = variable.value + group_results = [] + for condition in sub_conditions: + key = FileAttribute(condition.key) + values = [file_manager.get_attr(file=file, attr=key) for file in files] + sub_group_results = [ + _evaluate_condition( + value=value, + operator=condition.comparison_operator, + expected=condition.value, + ) + for value in values + ] + # Determine the result based on the presence of "not" in the comparison operator + result = all(sub_group_results) if "not" in condition.comparison_operator else any(sub_group_results) + group_results.append(result) + return all(group_results) if operator == "and" else any(group_results) diff --git a/api/core/workflow/utils/variable_template_parser.py b/api/core/workflow/utils/variable_template_parser.py new file mode 100644 index 0000000000000000000000000000000000000000..1d8fb38ebf8237acae2ec55f1dc9238ad515739c --- /dev/null +++ b/api/core/workflow/utils/variable_template_parser.py @@ -0,0 +1,131 @@ +import re +from collections.abc import Mapping, Sequence +from typing import Any + +from core.workflow.entities.variable_entities import VariableSelector + +REGEX = re.compile(r"\{\{(#[a-zA-Z0-9_]{1,50}(\.[a-zA-Z_][a-zA-Z0-9_]{0,29}){1,10}#)\}\}") + +SELECTOR_PATTERN = re.compile(r"\{\{(#[a-zA-Z0-9_]{1,50}(?:\.[a-zA-Z_][a-zA-Z0-9_]{0,29}){1,10}#)\}\}") + + +def extract_selectors_from_template(template: str, /) -> Sequence[VariableSelector]: + parts = SELECTOR_PATTERN.split(template) + selectors = [] + for part in filter(lambda x: x, parts): + if "." in part and part[0] == "#" and part[-1] == "#": + selectors.append(VariableSelector(variable=f"{part}", value_selector=part[1:-1].split("."))) + return selectors + + +class VariableTemplateParser: + """ + !NOTE: Consider to use the new `segments` module instead of this class. + + A class for parsing and manipulating template variables in a string. + + Rules: + + 1. Template variables must be enclosed in `{{}}`. + 2. The template variable Key can only be: #node_id.var1.var2#. + 3. The template variable Key cannot contain new lines or spaces, and must comply with rule 2. + + Example usage: + + template = "Hello, {{#node_id.query.name#}}! Your age is {{#node_id.query.age#}}." + parser = VariableTemplateParser(template) + + # Extract template variable keys + variable_keys = parser.extract() + print(variable_keys) + # Output: ['#node_id.query.name#', '#node_id.query.age#'] + + # Extract variable selectors + variable_selectors = parser.extract_variable_selectors() + print(variable_selectors) + # Output: [VariableSelector(variable='#node_id.query.name#', value_selector=['node_id', 'query', 'name']), + # VariableSelector(variable='#node_id.query.age#', value_selector=['node_id', 'query', 'age'])] + + # Format the template string + inputs = {'#node_id.query.name#': 'John', '#node_id.query.age#': 25}} + formatted_string = parser.format(inputs) + print(formatted_string) + # Output: "Hello, John! Your age is 25." + """ + + def __init__(self, template: str): + self.template = template + self.variable_keys = self.extract() + + def extract(self) -> list: + """ + Extracts all the template variable keys from the template string. + + Returns: + A list of template variable keys. + """ + # Regular expression to match the template rules + matches = re.findall(REGEX, self.template) + + first_group_matches = [match[0] for match in matches] + + return list(set(first_group_matches)) + + def extract_variable_selectors(self) -> list[VariableSelector]: + """ + Extracts the variable selectors from the template variable keys. + + Returns: + A list of VariableSelector objects representing the variable selectors. + """ + variable_selectors = [] + for variable_key in self.variable_keys: + remove_hash = variable_key.replace("#", "") + split_result = remove_hash.split(".") + if len(split_result) < 2: + continue + + variable_selectors.append(VariableSelector(variable=variable_key, value_selector=split_result)) + + return variable_selectors + + def format(self, inputs: Mapping[str, Any]) -> str: + """ + Formats the template string by replacing the template variables with their corresponding values. + + Args: + inputs: A dictionary containing the values for the template variables. + remove_template_variables: A boolean indicating whether to remove the template variables from the output. + + Returns: + The formatted string with template variables replaced by their values. + """ + + def replacer(match): + key = match.group(1) + value = inputs.get(key, match.group(0)) # return original matched string if key not found + + if value is None: + value = "" + # convert the value to string + if isinstance(value, list | dict | bool | int | float): + value = str(value) + + # remove template variables if required + return VariableTemplateParser.remove_template_variables(value) + + prompt = re.sub(REGEX, replacer, self.template) + return re.sub(r"<\|.*?\|>", "", prompt) + + @classmethod + def remove_template_variables(cls, text: str): + """ + Removes the template variables from the given text. + + Args: + text: The text from which to remove the template variables. + + Returns: + The text with template variables removed. + """ + return re.sub(REGEX, r"{\1}", text) diff --git a/api/core/workflow/workflow_engine_manager.py b/api/core/workflow/workflow_engine_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/api/core/workflow/workflow_entry.py b/api/core/workflow/workflow_entry.py new file mode 100644 index 0000000000000000000000000000000000000000..f622d0b2d01f28ee97937b19153f84aba7660225 --- /dev/null +++ b/api/core/workflow/workflow_entry.py @@ -0,0 +1,266 @@ +import logging +import time +import uuid +from collections.abc import Generator, Mapping, Sequence +from typing import Any, Optional + +from configs import dify_config +from core.app.apps.base_app_queue_manager import GenerateTaskStoppedError +from core.app.entities.app_invoke_entities import InvokeFrom +from core.file.models import File +from core.workflow.callbacks import WorkflowCallback +from core.workflow.entities.variable_pool import VariablePool +from core.workflow.errors import WorkflowNodeRunFailedError +from core.workflow.graph_engine.entities.event import GraphEngineEvent, GraphRunFailedEvent, InNodeEvent +from core.workflow.graph_engine.entities.graph import Graph +from core.workflow.graph_engine.entities.graph_init_params import GraphInitParams +from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState +from core.workflow.graph_engine.graph_engine import GraphEngine +from core.workflow.nodes import NodeType +from core.workflow.nodes.base import BaseNode +from core.workflow.nodes.event import NodeEvent +from core.workflow.nodes.node_mapping import NODE_TYPE_CLASSES_MAPPING +from factories import file_factory +from models.enums import UserFrom +from models.workflow import ( + Workflow, + WorkflowType, +) + +logger = logging.getLogger(__name__) + + +class WorkflowEntry: + def __init__( + self, + tenant_id: str, + app_id: str, + workflow_id: str, + workflow_type: WorkflowType, + graph_config: Mapping[str, Any], + graph: Graph, + user_id: str, + user_from: UserFrom, + invoke_from: InvokeFrom, + call_depth: int, + variable_pool: VariablePool, + thread_pool_id: Optional[str] = None, + ) -> None: + """ + Init workflow entry + :param tenant_id: tenant id + :param app_id: app id + :param workflow_id: workflow id + :param workflow_type: workflow type + :param graph_config: workflow graph config + :param graph: workflow graph + :param user_id: user id + :param user_from: user from + :param invoke_from: invoke from + :param call_depth: call depth + :param variable_pool: variable pool + :param thread_pool_id: thread pool id + """ + # check call depth + workflow_call_max_depth = dify_config.WORKFLOW_CALL_MAX_DEPTH + if call_depth > workflow_call_max_depth: + raise ValueError("Max workflow call depth {} reached.".format(workflow_call_max_depth)) + + # init workflow run state + self.graph_engine = GraphEngine( + tenant_id=tenant_id, + app_id=app_id, + workflow_type=workflow_type, + workflow_id=workflow_id, + user_id=user_id, + user_from=user_from, + invoke_from=invoke_from, + call_depth=call_depth, + graph=graph, + graph_config=graph_config, + variable_pool=variable_pool, + max_execution_steps=dify_config.WORKFLOW_MAX_EXECUTION_STEPS, + max_execution_time=dify_config.WORKFLOW_MAX_EXECUTION_TIME, + thread_pool_id=thread_pool_id, + ) + + def run( + self, + *, + callbacks: Sequence[WorkflowCallback], + ) -> Generator[GraphEngineEvent, None, None]: + """ + :param callbacks: workflow callbacks + """ + graph_engine = self.graph_engine + + try: + # run workflow + generator = graph_engine.run() + for event in generator: + if callbacks: + for callback in callbacks: + callback.on_event(event=event) + yield event + except GenerateTaskStoppedError: + pass + except Exception as e: + logger.exception("Unknown Error when workflow entry running") + if callbacks: + for callback in callbacks: + callback.on_event(event=GraphRunFailedEvent(error=str(e))) + return + + @classmethod + def single_step_run( + cls, + *, + workflow: Workflow, + node_id: str, + user_id: str, + user_inputs: dict, + ) -> tuple[BaseNode, Generator[NodeEvent | InNodeEvent, None, None]]: + """ + Single step run workflow node + :param workflow: Workflow instance + :param node_id: node id + :param user_id: user id + :param user_inputs: user inputs + :return: + """ + # fetch node info from workflow graph + workflow_graph = workflow.graph_dict + if not workflow_graph: + raise ValueError("workflow graph not found") + + nodes = workflow_graph.get("nodes") + if not nodes: + raise ValueError("nodes not found in workflow graph") + + # fetch node config from node id + try: + node_config = next(filter(lambda node: node["id"] == node_id, nodes)) + except StopIteration: + raise ValueError("node id not found in workflow graph") + + # Get node class + node_type = NodeType(node_config.get("data", {}).get("type")) + node_version = node_config.get("data", {}).get("version", "1") + node_cls = NODE_TYPE_CLASSES_MAPPING[node_type][node_version] + + # init variable pool + variable_pool = VariablePool(environment_variables=workflow.environment_variables) + + # init graph + graph = Graph.init(graph_config=workflow.graph_dict) + + # init workflow run state + node_instance = node_cls( + id=str(uuid.uuid4()), + config=node_config, + graph_init_params=GraphInitParams( + tenant_id=workflow.tenant_id, + app_id=workflow.app_id, + workflow_type=WorkflowType.value_of(workflow.type), + workflow_id=workflow.id, + graph_config=workflow.graph_dict, + user_id=user_id, + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + call_depth=0, + ), + graph=graph, + graph_runtime_state=GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()), + ) + + try: + # variable selector to variable mapping + variable_mapping = node_cls.extract_variable_selector_to_variable_mapping( + graph_config=workflow.graph_dict, config=node_config + ) + except NotImplementedError: + variable_mapping = {} + + cls.mapping_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id=workflow.tenant_id, + ) + try: + # run node + generator = node_instance.run() + except Exception as e: + raise WorkflowNodeRunFailedError(node_instance=node_instance, error=str(e)) + return node_instance, generator + + @staticmethod + def handle_special_values(value: Optional[Mapping[str, Any]]) -> Mapping[str, Any] | None: + result = WorkflowEntry._handle_special_values(value) + return result if isinstance(result, Mapping) or result is None else dict(result) + + @staticmethod + def _handle_special_values(value: Any) -> Any: + if value is None: + return value + if isinstance(value, dict): + res = {} + for k, v in value.items(): + res[k] = WorkflowEntry._handle_special_values(v) + return res + if isinstance(value, list): + res_list = [] + for item in value: + res_list.append(WorkflowEntry._handle_special_values(item)) + return res_list + if isinstance(value, File): + return value.to_dict() + return value + + @classmethod + def mapping_user_inputs_to_variable_pool( + cls, + *, + variable_mapping: Mapping[str, Sequence[str]], + user_inputs: dict, + variable_pool: VariablePool, + tenant_id: str, + ) -> None: + for node_variable, variable_selector in variable_mapping.items(): + # fetch node id and variable key from node_variable + node_variable_list = node_variable.split(".") + if len(node_variable_list) < 1: + raise ValueError(f"Invalid node variable {node_variable}") + + node_variable_key = ".".join(node_variable_list[1:]) + + if (node_variable_key not in user_inputs and node_variable not in user_inputs) and not variable_pool.get( + variable_selector + ): + raise ValueError(f"Variable key {node_variable} not found in user inputs.") + + # environment variable already exist in variable pool, not from user inputs + if variable_pool.get(variable_selector): + continue + + # fetch variable node id from variable selector + variable_node_id = variable_selector[0] + variable_key_list = variable_selector[1:] + variable_key_list = list(variable_key_list) + + # get input value + input_value = user_inputs.get(node_variable) + if not input_value: + input_value = user_inputs.get(node_variable_key) + + if isinstance(input_value, dict) and "type" in input_value and "transfer_method" in input_value: + input_value = file_factory.build_from_mapping(mapping=input_value, tenant_id=tenant_id) + if ( + isinstance(input_value, list) + and all(isinstance(item, dict) for item in input_value) + and all("type" in item and "transfer_method" in item for item in input_value) + ): + input_value = file_factory.build_from_mappings(mappings=input_value, tenant_id=tenant_id) + + # append variable and value to variable pool + variable_pool.add([variable_node_id] + variable_key_list, input_value)