knowledge.py
from flask import request from flask_restful import Resource, reqparse from bedrock.knowledge_service import ExternalDatasetService class BedrockRetrievalApi(Resource): # url : <your-endpoint>/retrieval def post(self): parser = reqparse.RequestParser() parser.add_argument("retrieval_setting", nullable=False, required=True, type=dict, location="json") parser.add_argument("query", nullable=False, required=True, type=str,) parser.add_argument("knowledge_id", nullable=False, required=True, type=str) args = parser.parse_args() # Authorization check auth_header = request.headers.get("Authorization") if " " not in auth_header: return { "error_code": 1001, "error_msg": "Invalid Authorization header format. Expected 'Bearer <api-key>' format." }, 403 auth_scheme, auth_token = auth_header.split(None, 1) auth_scheme = auth_scheme.lower() if auth_scheme != "bearer": return { "error_code": 1001, "error_msg": "Invalid Authorization header format. Expected 'Bearer <api-key>' format." }, 403 if auth_token: # process your authorization logic here pass # Call the knowledge retrieval service result = ExternalDatasetService.knowledge_retrieval( args["retrieval_setting"], args["query"], args["knowledge_id"] ) return result, 200
knowledge_service.py
import boto3 class ExternalDatasetService: @staticmethod def knowledge_retrieval(retrieval_setting: dict, query: str, knowledge_id: str): # get bedrock client client = boto3.client( "bedrock-agent-runtime", aws_secret_access_key="AWS_SECRET_ACCESS_KEY", aws_access_key_id="AWS_ACCESS_KEY_ID", # example: us-east-1 region_name="AWS_REGION_NAME", ) # fetch external knowledge retrieval response = client.retrieve( knowledgeBaseId=knowledge_id, retrievalConfiguration={ "vectorSearchConfiguration": {"numberOfResults": retrieval_setting.get("top_k"), "overrideSearchType": "HYBRID"} }, retrievalQuery={"text": query}, ) # parse response results = [] if response.get("ResponseMetadata") and response.get("ResponseMetadata").get("HTTPStatusCode") == 200: if response.get("retrievalResults"): retrieval_results = response.get("retrievalResults") for retrieval_result in retrieval_results: # filter out results with score less than threshold if retrieval_result.get("score") < retrieval_setting.get("score_threshold", .0): continue result = { "metadata": retrieval_result.get("metadata"), "score": retrieval_result.get("score"), "title": retrieval_result.get("metadata").get("x-amz-bedrock-kb-source-uri"), "content": retrieval_result.get("content").get("text"), } results.append(result) return { "records": results }
api-endpoint/retrieval
Was this page helpful?