feat: add inner api key

This commit is contained in:
Yeuoly 2024-09-20 13:32:11 +08:00
parent 661392eaef
commit 73ce692e24
No known key found for this signature in database
GPG Key ID: A66E7E320FB19F61
4 changed files with 26 additions and 17 deletions

View File

@ -278,10 +278,6 @@ WORKFLOW_CALL_MAX_DEPTH=5
APP_MAX_EXECUTION_TIME=1200 APP_MAX_EXECUTION_TIME=1200
APP_MAX_ACTIVE_REQUESTS=0 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 configuration
CELERY_BEAT_SCHEDULER_TIME=1 CELERY_BEAT_SCHEDULER_TIME=1
@ -293,3 +289,8 @@ POSITION_TOOL_EXCLUDES=
POSITION_PROVIDER_PINS= POSITION_PROVIDER_PINS=
POSITION_PROVIDER_INCLUDES= POSITION_PROVIDER_INCLUDES=
POSITION_PROVIDER_EXCLUDES= 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

View File

@ -115,14 +115,19 @@ class PluginConfig(BaseSettings):
""" """
Plugin configs Plugin configs
""" """
PLUGIN_INNER_API_URL: str = Field( PLUGIN_API_URL: str = Field(
description='Plugin inner API URL', description='Plugin API URL',
default='http://plugin:8194', default='http://plugin:5002',
) )
PLUGIN_INNER_API_KEY: str = Field( PLUGIN_API_KEY: str = Field(
description='Plugin inner API key', description='Plugin API key',
default='dify-inner-api-key', default='plugin-api-key',
)
INNER_API_KEY_FOR_PLUGIN: str = Field(
description='Inner api key for plugin',
default='inner-api-key'
) )

View File

@ -18,7 +18,7 @@ def enterprise_inner_api_only(view):
# get header 'X-Inner-Api-Key' # get header 'X-Inner-Api-Key'
inner_api_key = request.headers.get("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) abort(401)
return view(*args, **kwargs) return view(*args, **kwargs)
@ -67,14 +67,14 @@ def enterprise_inner_api_user_auth(view):
def plugin_inner_api_only(view): def plugin_inner_api_only(view):
@wraps(view) @wraps(view)
def decorated(*args, **kwargs): def decorated(*args, **kwargs):
if not dify_config.PLUGIN_INNER_API_KEY: if not dify_config.PLUGIN_API_KEY:
abort(404) abort(404)
# get header 'X-Inner-Api-Key' # get header 'X-Inner-Api-Key'
inner_api_key = request.headers.get('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: if not inner_api_key or inner_api_key != dify_config.INNER_API_KEY_FOR_PLUGIN:
abort(404) abort(404)
return view(*args, **kwargs) return view(*args, **kwargs)
return decorated return decorated

View File

@ -49,7 +49,7 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool):
retrieval_thread = threading.Thread( retrieval_thread = threading.Thread(
target=self._retriever, target=self._retriever,
kwargs={ kwargs={
"flask_app": current_app._get_current_object(), "flask_app": current_app._get_current_object(), # type: ignore
"dataset_id": dataset_id, "dataset_id": dataset_id,
"query": query, "query": query,
"all_documents": all_documents, "all_documents": all_documents,
@ -77,11 +77,12 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool):
document_score_list = {} document_score_list = {}
for item in all_documents: for item in all_documents:
assert item.metadata
if item.metadata.get("score"): if item.metadata.get("score"):
document_score_list[item.metadata["doc_id"]] = item.metadata["score"] document_score_list[item.metadata["doc_id"]] = item.metadata["score"]
document_context_list = [] 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( segments = DocumentSegment.query.filter(
DocumentSegment.dataset_id.in_(self.dataset_ids), DocumentSegment.dataset_id.in_(self.dataset_ids),
DocumentSegment.completed_at.isnot(None), DocumentSegment.completed_at.isnot(None),
@ -140,6 +141,8 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool):
return str("\n".join(document_context_list)) return str("\n".join(document_context_list))
raise RuntimeError("not segments found")
def _retriever( def _retriever(
self, self,
flask_app: Flask, flask_app: Flask,