Lab 04: 관리형 KB를 코드로 호출

• 콘솔에서 만든 KB를 bedrock-agent-runtime 클라이언트로 코드에서 호출할 수 있습니다
• 모델 호출(bedrock-runtime)과 KB 호출(bedrock-agent-runtime)의 클라이언트 차이를 설명할 수 있습니다
• retrieve(검색만)와 retrieve_and_generate(검색+생성)의 역할 차이를 이해합니다
• retrieve_and_generate의 modelArn에 인퍼런스 프로파일 ARN을 넘기는 이유를 압니다
• 본인 KB ID를 환경변수 또는 설정으로 주입하여 코드를 실행할 수 있습니다

오전 콘솔 실습에서 우리는 사실 두 종류의 서로 다른 백엔드를 다뤘습니다. 하나는 모델 자체(Claude, Titan)이고, 다른 하나는 그 모델 위에 검색 기능을 얹은 관리형 RAG 서비스(Knowledge Base)입니다. boto3에서도 이 둘은 별도의 클라이언트로 갈라집니다. 이 구분을 놓치면 “함수가 없다”는 오류로 한참 헤매게 되므로, 가장 먼저 확실히 잡고 갑니다.

챗봇 과정에서 쓴 것이 bedrock-runtime이었습니다. 그 클라이언트에는 retrieve 메서드가 없습니다. KB를 호출하려면 반드시 bedrock-agent-runtime을 새로 만들어야 합니다. (참고로 bedrock이라는 또 다른 클라이언트도 있는데, 그것은 KB를 “생성/삭제/관리”하는 컨트롤 플레인용이며, 호출용인 두 런타임과 또 다릅니다.)

넥스트클라우드 제공 모듈 config.py는 이 두 클라이언트를 각각 만들어주는 헬퍼를 이미 제공합니다. 학습자는 이 헬퍼만 가져다 쓰면 됩니다.

def get_runtime():
    """모델 호출용 bedrock-runtime 클라이언트."""
    return boto3.client("bedrock-runtime", region_name=REGION)


def get_agent_runtime():
    """관리형 KB 호출용 bedrock-agent-runtime 클라이언트."""
    return boto3.client("bedrock-agent-runtime", region_name=REGION)

• 두 헬퍼는 같은 리전(us-east-1)을 사용하지만, 만들어내는 클라이언트의 종류가 다릅니다.
• 이번 Lab에서 KB를 호출할 때는 get_agent_runtime()만 사용합니다.

graph LR
  Q["질문"] --> AR["bedrock-agent-runtime"]
  AR -->|retrieve| KB["Knowledge Base<br/>(벡터 검색)"]
  AR -->|retrieve_and_generate| KB
  KB -.->|검색+생성 시| BR["bedrock-runtime<br/>(모델 생성)"]
  KB --> R["검색 결과 / 답변"]

오전 콘솔 실습 마지막에 각자 본인 KB의 ID를 메모해 두었을 것입니다. 코드가 호출할 대상이 바로 그 ID입니다. KB ID는 사람마다 다르므로 코드에 하드코딩하지 않고, 환경변수로 주입하거나 설정 파일에서 읽어옵니다.

config.py는 KB ID를 환경변수 KB_ID에서 먼저 찾고, 없으면 자리표시자를 사용합니다.

REGION = "us-east-1"

# 생성(답변)용 모델 — 챗봇 과정과 동일한 최신 Haiku 4.5 (인퍼런스 프로파일).
MODEL_HAIKU = "us.anthropic.claude-haiku-4-5-20251001-v1:0"
MODEL_NOVA = "us.amazon.nova-lite-v1:0"

# 임베딩 모델 — Titan Text Embeddings V2 (출력 차원 1024).
EMBED_MODEL = "amazon.titan-embed-text-v2:0"

# 콘솔에서 만든 관리형 Knowledge Base 의 ID.
KB_ID = os.environ.get("KB_ID", "<여기에-본인-KB-ID>")

• KB_ID가 비어 있거나 오타이면 호출 시 ResourceNotFoundException이 발생합니다. 본인 KB가 콘솔에서 Available 상태인지, 리전이 us-east-1인지 함께 확인하세요.
• MODEL_HAIKU가 오후 답변 생성에 쓰이는 주력 모델입니다. 챗봇 과정과 동일한 Haiku 4.5(인퍼런스 프로파일 형식의 모델 ID)입니다.
• 다른 제공사 모델로 바꿔 보려면 MODEL_NOVA(Amazon Nova)로 스왑할 수 있습니다.

실행할 때는 다음처럼 환경변수로 본인 KB ID를 한 번에 넘기는 것이 가장 간편합니다.

# 본인 KB ID 를 환경변수로 주입하여 실행
KB_ID=XXXXXXXX python lab4_kb_retrieve.py

retrieve_and_generate는 검색만 하는 것이 아니라 모델을 불러 답변까지 생성합니다. 그래서 “어떤 모델로 답을 쓸지”를 알려줘야 하는데, 이때 넘기는 값이 modelArn입니다. 여기서 흔히 막히는 함정이 있습니다. 단순 모델 ID 문자열을 그대로 넣으면 실패합니다.

우리가 쓰는 Haiku 4.5 같은 최신 모델은 인퍼런스 프로파일(inference profile) 형태로 호출해야 합니다. 인퍼런스 프로파일은 여러 리전에 걸쳐 요청을 분산해주는 교차 리전 호출 단위로, us.anthropic...처럼 us. 접두사가 붙은 ID가 그 표시입니다. retrieve_and_generate의 modelArn에는 이 인퍼런스 프로파일을 가리키는 전체 ARN을 넘겨야 합니다.

config.py의 model_arn 헬퍼가 계정 ID를 조회해서 올바른 ARN을 자동으로 조립해줍니다.

def model_arn(model_id: str = MODEL_HAIKU) -> str:
    """retrieve_and_generate 등에 넘길 모델 ARN(인퍼런스 프로파일 ARN) 생성."""
    acct = boto3.client("sts", region_name=REGION).get_caller_identity()["Account"]
    return f"arn:aws:bedrock:{REGION}:{acct}:inference-profile/{model_id}"

• sts().get_caller_identity()로 현재 자격 증명의 계정 ID를 가져옵니다. ARN에는 계정 ID가 들어가므로 직접 호출 계정에 맞춰 만들어집니다.
• 결과는 arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.anthropic.claude-haiku-4-5-20251001-v1:0 형태입니다. 끝부분이 inference-profile/이라는 점에 주목하세요. 이것이 “프로파일 ARN”이라는 표시입니다.
• 만약 modelArn에 bare 모델 ID(us.anthropic...만)나 foundation-model/... ARN을 넣으면 검증 오류가 납니다. 반드시 model_arn()이 만들어주는 프로파일 ARN을 쓰세요.

retrieve는 콘솔 Test의 “Retrieval only” 토글에 해당합니다. 질문을 벡터로 바꿔 KB 안에서 가장 비슷한 청크 몇 개를 찾아 그대로 돌려줄 뿐, 모델로 답변을 생성하지는 않습니다. “검색 단계만” 떼어내 무엇이 검색됐는지, 점수가 얼마인지 눈으로 확인할 때 유용합니다.

from config import get_agent_runtime, KB_ID, model_arn, MODEL_HAIKU

agent_rt = get_agent_runtime()


def retrieve(query: str, top_k: int = 3) -> list[dict]:
    """KB 에서 질문과 관련된 청크를 검색만 한다(생성 없음)."""
    r = agent_rt.retrieve(
        knowledgeBaseId=KB_ID,
        retrievalQuery={"text": query},
        retrievalConfiguration={"vectorSearchConfiguration": {"numberOfResults": top_k}},
    )
    return r["retrievalResults"]

한 줄씩 살펴봅니다.

• agent_rt = get_agent_runtime() — KB 호출 전용 클라이언트를 모듈 로드 시점에 한 번 만들어 둡니다.
• knowledgeBaseId=KB_ID — 어떤 KB에 물어볼지 지정합니다. 본인이 주입한 KB ID가 여기로 들어갑니다.
• retrievalQuery={"text": query} — 검색할 질문 텍스트입니다. 이 텍스트가 내부에서 임베딩되어 벡터 검색에 쓰입니다.
• vectorSearchConfiguration: {numberOfResults: top_k} — 가장 유사한 청크를 몇 개(top-k) 가져올지입니다. 기본 3개로 두었습니다.
• 반환값 r["retrievalResults"]는 청크들의 리스트입니다. 각 항목에는 본문 텍스트(content.text), 유사도 점수(score), 출처(location) 등이 들어 있습니다.

top_k를 1로 줄이면 어떤 일이 생길까요? 답에 필요한 정보가 두 번째 청크에 있었다면 그 청크가 빠져 답이 부실해집니다. 검색이 RAG 품질의 절반이라는 말이 여기서 체감됩니다.

retrieve_and_generate는 콘솔 Test의 “Retrieval and generation” 토글입니다. 검색(retrieve)부터 모델 답변 생성까지 Bedrock이 한 번의 API 호출로 모두 처리합니다. 우리가 챗봇 과정에서 직접 했던 “검색 결과를 프롬프트에 끼워 모델에 넣기”를 AWS가 대신 해주는 셈입니다.

def retrieve_and_generate(query: str) -> tuple[str, list]:
    """KB 검색 + 답변 생성을 Bedrock 이 한 번에 처리한다. (답변, 인용목록) 반환."""
    r = agent_rt.retrieve_and_generate(
        input={"text": query},
        retrieveAndGenerateConfiguration={
            "type": "KNOWLEDGE_BASE",
            "knowledgeBaseConfiguration": {
                "knowledgeBaseId": KB_ID,
                "modelArn": model_arn(MODEL_HAIKU),  # 인퍼런스 프로파일 ARN
            },
        },
    )
    return r["output"]["text"], r.get("citations", [])

한 줄씩 살펴봅니다.

• input={"text": query} — 사용자 질문입니다. retrieve의 retrievalQuery와 키 이름이 다르다는 점에 주의하세요(여기는 input).
• type: "KNOWLEDGE_BASE" — 이 호출이 KB 기반 RAG임을 명시합니다.
• knowledgeBaseId: KB_ID — retrieve와 마찬가지로 대상 KB를 지정합니다.
• modelArn: model_arn(MODEL_HAIKU) — 검색 결과로 답을 생성할 모델입니다. 앞에서 본 대로 반드시 인퍼런스 프로파일 ARN이어야 하므로 model_arn() 헬퍼를 통해 넘깁니다.
• 반환값 r["output"]["text"]가 최종 답변 텍스트입니다.
• r.get("citations", [])는 인용(citation) 목록입니다. 답변의 어느 부분이 어느 문서 청크에 근거했는지를 담고 있어, 콘솔 Test에서 “출처 클릭”으로 보던 그 근거를 코드에서도 그대로 받을 수 있습니다.

retrieve와 retrieve_and_generate를 나란히 보면, 전자는 “검색 결과(재료)“만 주고 후자는 “완성된 답변+근거”까지 준다는 차이가 분명합니다. 콘솔 Test의 두 토글이 곧 이 두 API의 GUI였던 것입니다.

lab4_kb_retrieve.py의 메인 블록은 같은 질문 하나로 두 함수를 차례로 호출하여, 검색만 했을 때와 검색+생성을 했을 때의 결과를 비교해 보여줍니다.

if __name__ == "__main__":
    q = "이 문서의 핵심 주제를 한 문장으로 알려줘."

    print("== retrieve (검색만) ==")
    for i, h in enumerate(retrieve(q)):
        score = h.get("score")
        snippet = h["content"]["text"][:80].replace("\n", " ")
        print(f"  [{i}] score={score:.3f}  {snippet}")

    print("\n== retrieve_and_generate (검색+생성) ==")
    answer, citations = retrieve_and_generate(q)
    print("  답변:", answer)
    print(f"  인용 {len(citations)}개")

• retrieve(q) 결과는 청크 리스트이므로 하나씩 돌며 점수(score)와 본문 앞 80자(snippet)를 출력합니다. 점수가 높을수록 질문과 의미적으로 가깝다는 뜻입니다.
• retrieve_and_generate(q)는 완성된 답변과 인용 목록을 돌려주므로, 답변 본문과 인용 개수를 출력합니다.

이제 본인 KB ID를 넣어 실행해 봅니다.

# RAG 실습 폴더로 이동 (clone 받은 레포)
cd ~/nxt-workshop-code/bedrock-rag/starter

# 본인 KB ID 를 환경변수로 주입하여 실행
KB_ID=XXXXXXXX python lab4_kb_retrieve.py

출력에서 다음을 확인하세요.

• retrieve 결과의 각 청크 점수가 합리적인지(질문과 동떨어진 청크가 1순위로 올라오지는 않는지).
• retrieve_and_generate의 답변이 검색된 청크 내용에 근거하고 있는지, 인용 개수가 0이 아닌지.

콘솔에서 마우스로 토글하던 두 모드를, 이제 코드로 똑같이 재현했습니다. “콘솔 Test = 이 API들의 GUI였다”는 사실을 직접 체감하는 순간입니다.

기억할 세 가지입니다.

• 클라이언트가 다르다 — 모델 호출은 bedrock-runtime, KB 호출은 bedrock-agent-runtime.
• 두 함수의 역할이 다르다 — retrieve는 검색 재료만, retrieve_and_generate는 답변+근거까지.
• modelArn은 프로파일 ARN — 최신 모델은 인퍼런스 프로파일 ARN으로만 호출되므로 model_arn() 헬퍼를 통한다.

관리형 KB는 청킹·임베딩·벡터검색·생성을 모두 AWS가 대신 처리해주기에 코드가 이렇게 짧습니다. 다음 Lab에서는 이 “마법”을 직접 해체합니다.

• numberOfResults를 1 vs 10으로 바꾸면 retrieve_and_generate 답이 어떻게 달라질까요? 많을수록 항상 좋을까요?
• retrieve의 score는 무엇을 재는 값일까요? 점수가 낮은 청크가 섞이면?