From 73ce692e24153f76652e306711ea3af6b3691cf4 Mon Sep 17 00:00:00 2001 From: Yeuoly Date: Fri, 20 Sep 2024 13:32:11 +0800 Subject: [PATCH] feat: add inner api key --- api/.env.example | 9 +++++---- api/configs/feature/__init__.py | 17 +++++++++++------ api/controllers/inner_api/wraps.py | 10 +++++----- .../dataset_multi_retriever_tool.py | 7 +++++-- 4 files changed, 26 insertions(+), 17 deletions(-) diff --git a/api/.env.example b/api/.env.example index 21940402f4..3cff88c400 100644 --- a/api/.env.example +++ b/api/.env.example @@ -278,10 +278,6 @@ WORKFLOW_CALL_MAX_DEPTH=5 APP_MAX_EXECUTION_TIME=1200 APP_MAX_ACTIVE_REQUESTS=0 -# Plugin configuration -PLUGIN_INNER_API_URL=http://127.0.0.1:5002 -PLUGIN_INNER_API_KEY=lYkiYYT6owG+71oLerGzA7GXCgOT++6ovaezWAjpCjf+Sjc3ZtU+qUEi - # Celery beat configuration CELERY_BEAT_SCHEDULER_TIME=1 @@ -293,3 +289,8 @@ POSITION_TOOL_EXCLUDES= POSITION_PROVIDER_PINS= POSITION_PROVIDER_INCLUDES= POSITION_PROVIDER_EXCLUDES= + +# Plugin configuration +PLUGIN_API_KEY=lYkiYYT6owG+71oLerGzA7GXCgOT++6ovaezWAjpCjf+Sjc3ZtU+qUEi+vRjI/+XbV1AaFy691iy+kGDv2Jvy0/eAh8Y1 +PLUGIN_API_URL=http://127.0.0.1:5002 +INNER_API_KEY=QaHbTe77CtuXmsfyhR7+vRjI/+XbV1AaFy691iy+kGDv2Jvy0/eAh8Y1 \ No newline at end of file diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py index 4e1dfe73ad..30b97749ee 100644 --- a/api/configs/feature/__init__.py +++ b/api/configs/feature/__init__.py @@ -115,14 +115,19 @@ class PluginConfig(BaseSettings): """ Plugin configs """ - PLUGIN_INNER_API_URL: str = Field( - description='Plugin inner API URL', - default='http://plugin:8194', + PLUGIN_API_URL: str = Field( + description='Plugin API URL', + default='http://plugin:5002', ) - PLUGIN_INNER_API_KEY: str = Field( - description='Plugin inner API key', - default='dify-inner-api-key', + PLUGIN_API_KEY: str = Field( + description='Plugin API key', + default='plugin-api-key', + ) + + INNER_API_KEY_FOR_PLUGIN: str = Field( + description='Inner api key for plugin', + default='inner-api-key' ) diff --git a/api/controllers/inner_api/wraps.py b/api/controllers/inner_api/wraps.py index 3e184515d3..b543b5e164 100644 --- a/api/controllers/inner_api/wraps.py +++ b/api/controllers/inner_api/wraps.py @@ -18,7 +18,7 @@ def enterprise_inner_api_only(view): # get header 'X-Inner-Api-Key' inner_api_key = request.headers.get("X-Inner-Api-Key") - if not inner_api_key or inner_api_key != dify_config.INNER_API_KEY: + if not inner_api_key or inner_api_key != dify_config.INNER_API_KEY_FOR_PLUGIN: abort(401) return view(*args, **kwargs) @@ -67,14 +67,14 @@ def enterprise_inner_api_user_auth(view): def plugin_inner_api_only(view): @wraps(view) def decorated(*args, **kwargs): - if not dify_config.PLUGIN_INNER_API_KEY: + if not dify_config.PLUGIN_API_KEY: abort(404) # get header 'X-Inner-Api-Key' - inner_api_key = request.headers.get('X-Inner-Api-Key') - if not inner_api_key or inner_api_key != dify_config.PLUGIN_INNER_API_KEY: + inner_api_key = request.headers.get("X-Inner-Api-Key") + if not inner_api_key or inner_api_key != dify_config.INNER_API_KEY_FOR_PLUGIN: abort(404) return view(*args, **kwargs) - return decorated \ No newline at end of file + return decorated diff --git a/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py b/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py index 408dd43d22..d3a8752d9a 100644 --- a/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py +++ b/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py @@ -49,7 +49,7 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool): retrieval_thread = threading.Thread( target=self._retriever, kwargs={ - "flask_app": current_app._get_current_object(), + "flask_app": current_app._get_current_object(), # type: ignore "dataset_id": dataset_id, "query": query, "all_documents": all_documents, @@ -77,11 +77,12 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool): document_score_list = {} for item in all_documents: + assert item.metadata if 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] + 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), @@ -140,6 +141,8 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool): return str("\n".join(document_context_list)) + raise RuntimeError("not segments found") + def _retriever( self, flask_app: Flask,