원제: A-RAG: Scaling Agentic Retrieval-Augmented Generation via
Hierarchical Retrieval Interfaces

사전 세팅

arag/scripts/build_index.py

function: semantic search용 embedding index를 미리 만드는 스크립트

• chunks.json을 입력받음
• load_chunks: 파일을 입력받아서 dictionary형태로 반환
  • {"id": ..., "text": ...}형태로 전부 반환
• split_sentences: chunk를 문장 단위(.!?\n)로 split
  • 이때 빈 문자열과 길이가 10자 이하인 짧은 조각은 제거함

build_index

• load_chunk로 청크단위로 만듦
• split_sentences를 이용해 청크를 문장 단위로 나누고 나눠진 청크에 맞춰 chunk id를 sentence_to_chunk에 순서대로 매핑
• embedding model: 기본값은 all-MiniLM-L6-v2
• output: pkl파일로 sentences, embeddings, chunk id, chunk, model name을 반환
  • pkl: pickle 라이브러리의 기능으로 데이터를 binary 형태로 변환

arag/src/arag/tools/registry.py

• ToolRegistry: 도구를 등록해두는 역할
  • LLM이 tool(keyword, semantic, read_chunk)를 문자열로 선택하게 만듦
  • get_all_schemas: 어떤 tool들을 사용할 수 있는지를 알리는 역할

arag/scripts/batch_runner.py

_init_shared_tools

• KeywordSearchTool, ReadChunkTool을 ToolRegistry에 저장
• build_index에서 만든 pkl 파일을 이용해 embedding 모델 로드 후 SemanticSearchTool도 저장

_load_question: limit개 만큼의 query만 처리(테스트용)
_load_completed_qids: 이미 처리된 질문은 스킵

• _append_prediction: 이미 처리된 질문 리스트 제작

_create_agent: 위의 tool바탕 LLM 호출
_process_one: 한 질문에 대해서 처리
run: 이미 처리된 질문들 제외, pending questions에 대해서 처리

main tools

keyword_search

tiktoken

• LLM 사용을 위해 텍스트를 토큰화하고 토큰 수를 계산하거나 토큰을 다시 텍스트로 변환해주는 도구

실제 작동

• _load_chunks와 _split_sentences는 위와 동일
• name, get_schema: LLM에게 keyword_search에 대해서 전달

execute

• top k: 기본값 5, 최대 20
• chunk의 text와 입력받은 keyword를 모두 lower로 변환하고 count를 이용해서 chunk별로 해당 keyword가 얼마나 들어있는지를 확인
  • 최종 점수는 키워드 개수 $\times$ 키워드의 길이
• 만약 chunk 전체에 키워드가 한 개라도 포함되어 있다면(최종 점수가 0보다 크다면) 해당 키워드가 포함된 sentence를 찾음
  • 이후 scored_chunks에 chunk id와 점수, match된 sentence를 저장하고 sort를 이용해 top_chunk들을 찾음
  • 이때 토큰 관리를 위해 match된 sentence는 5개만 저장
• top_chunks에 대해 snippet을 생성
  • 정확히 그 단어가 들어간 문장만 따로 저장
  • 이때 앞뒤로 ...을 붙여서 snippet임을 알림
• 이후에 matched sentence를 리스트형태로 저장하고, 이를 sentences_text라는 하나의 텍스트로 결합한 뒤 token의 개수를 구해 저장합니다.
• return은 토큰의 개수와 top_chunk를 전달합니다.

semantic_search

앞쪽은 keyword_search와 거의 동일

• _load_index를 통해 임베딩을 가져옴

execute

• top k: 기본값 5, 최대 20
• query를 임베딩 + normalization
  • 이때 multi-thread에서도 안전하게 실행되도록 lock을 걺
• similarities: chunk별 임베딩값과 query 임베딩값을 dot product하여 cosine similarites 계산
• top_indices
  • argsort: similarities 바탕 오름차순 전환 $\rightarrow$ [::-1]: 내림차순으로 뒤집음 $\rightarrow$ top_k의 3배를 가져옴(노이즈가 존재하기에)
• 다시 청크화
  • top_indices들에 대해 sentences, 유사도, top_indices에서의 position을 각각 저장
  • 이후 동일한 청크는 같은 키값의 chunk_sentences에 저장
• chunk_sentences에서 청크별로 max similarity를 구함
  • 이후 chunk_scores에 chunk id, max similarity, sent(sentences, 각각 유사도, position)를 tuple형태로 저장
  • chunk_scores는 청크별로 max similarity를 기준으로 다시 sort
  • 이후 top_k를 뽑아냄
• k개에 대해서
  • 청크 내에서 유사도 기준이 아닌 원래 문장 순서를 기준으로 재정렬
  • 이후 snippet 생성
  • result_parts에 chunk ID, similarity, matched text를 저장
• 이후 최종 top_k개의 청크와 토큰 개수, 찾은 청크 수 반환

read_chunk

• keyword와 동일하게 청크를 dictionary형태로 변환

execute

• arguments
  • context: 현재 에이전트의 상태
  • chunk_ids: 여러 청크를 한 번에 읽기 위한 리스트
  • chunk_id: 하나의 청크만 읽을 때 사용하는 단일 ID -> backward compatibility
• chunk_ids $\rightarrow$ str형태로 변환(키값으로 쓰기 위함)
• 이미 읽은 청크인지 확인(is_chunk_read)
  • 출력은 하되 내용은 다시 보여주지 않음.
  • 이미 읽음 표시만 남김
  • continue로 이후 로직을 건너뜀
• 새로운 chunk인 경우
  • 실제 chunk 내용을 가져옴
  • 새로 읽은 chunk는 토큰 계산
  • 이후 읽음 상태로 기록(mark_chunk_as_read)
• 이후 too_result로 합치고 log(토큰 개수, 새로 읽은 청크 개수, 이미 읽은 청크 개수)와 함께 반환