dify/api/controllers/inner_api/knowledge/retrieval.py
盐粒 Yanli 0ea0647dd0
feat(agent): wire knowledge base retrieval into runtime (#37577)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-06-17 09:27:38 +00:00

111 lines
4.4 KiB
Python

"""Inner API endpoint for tenant-scoped knowledge retrieval.
This controller is a thin HTTP wrapper around
``services.knowledge_retrieval_inner_service.InnerKnowledgeRetrievalService``.
It intentionally keeps authorization simple: shared inner API key plus
tenant-scoped app/dataset validation in the service layer.
"""
from flask_restx import Resource
from pydantic import ValidationError
from controllers.common.schema import register_response_schema_models, register_schema_models
from controllers.inner_api import inner_api_ns
from controllers.inner_api.wraps import inner_api_only
from core.workflow.nodes.knowledge_retrieval import exc as retrieval_exc
from libs.exception import BaseHTTPException
from services.entities.knowledge_retrieval_inner import InnerKnowledgeRetrieveRequest, InnerKnowledgeRetrieveResponse
from services.errors.knowledge_retrieval import ExternalKnowledgeRetrievalError, InnerKnowledgeRetrievalServiceError
from services.knowledge_retrieval_inner_service import InnerKnowledgeRetrievalService
class InnerKnowledgeRetrievalHttpError(BaseHTTPException):
error_code = "knowledge_retrieve_failed"
description = "Knowledge retrieval failed."
code = 500
def __init__(
self,
*,
error_code: str | None = None,
description: str | None = None,
status_code: int | None = None,
) -> None:
if error_code is not None:
self.error_code = error_code
if description is not None:
self.description = description
if status_code is not None:
self.code = status_code
super().__init__(self.description)
register_schema_models(inner_api_ns, InnerKnowledgeRetrieveRequest)
register_response_schema_models(inner_api_ns, InnerKnowledgeRetrieveResponse)
@inner_api_ns.route("/knowledge/retrieve")
class InnerKnowledgeRetrieveApi(Resource):
"""Retrieve knowledge from one or more datasets within the caller tenant."""
@inner_api_only
@inner_api_ns.doc("inner_knowledge_retrieve")
@inner_api_ns.doc(description="Retrieve knowledge for trusted internal callers")
@inner_api_ns.expect(inner_api_ns.models[InnerKnowledgeRetrieveRequest.__name__])
@inner_api_ns.response(
200,
"Knowledge retrieved successfully",
inner_api_ns.models[InnerKnowledgeRetrieveResponse.__name__],
)
@inner_api_ns.doc(
responses={
400: "Invalid request body",
401: "Unauthorized - invalid inner API key",
403: "Caller tenant does not own the requested resource",
404: "App or dataset not found",
422: "Invalid retrieval configuration",
429: "Knowledge retrieval rate limited",
502: "External knowledge retrieval failed",
500: "Unexpected knowledge retrieval failure",
}
)
def post(self) -> dict[str, object]:
"""Validate the payload, run retrieval, and return workflow-style sources."""
try:
payload = InnerKnowledgeRetrieveRequest.model_validate(inner_api_ns.payload or {})
except ValidationError as exc:
raise InnerKnowledgeRetrievalHttpError(
error_code="invalid_request",
description=str(exc),
status_code=400,
) from exc
try:
response = InnerKnowledgeRetrievalService().retrieve(payload)
except InnerKnowledgeRetrievalServiceError as exc:
raise InnerKnowledgeRetrievalHttpError(
error_code=exc.error_code,
description=exc.description,
status_code=exc.status_code,
) from exc
except retrieval_exc.RateLimitExceededError as exc:
raise InnerKnowledgeRetrievalHttpError(
error_code="knowledge_rate_limited",
description=str(exc),
status_code=429,
) from exc
except ExternalKnowledgeRetrievalError as exc:
raise InnerKnowledgeRetrievalHttpError(
error_code="external_knowledge_failed",
description=str(exc),
status_code=502,
) from exc
except ValueError as exc:
raise InnerKnowledgeRetrievalHttpError(
error_code="retrieval_config_invalid",
description=str(exc),
status_code=422,
) from exc
return response.model_dump(mode="json", by_alias=True)