RAG Caching (/w Langchain & Langgraph)
LangGraph를 활용한 RAG 시스템의 챗봇을 과제로 진행하던 도중, 팀장님의 조언으로 동일한 질문에 대해 캐싱을 사용해 비용과 시간을 효율적으로 사용하라는 말씀을 들었습니다. 이에 따라 캐싱을 사용해 보다 더 효율적인 RAG 시스템을 구축하고자 하였습니다.
1. 문제상황
아래는 사용자의 같은 질문에 대해 소요된 시간과 토큰에 대한 이미지입니다.
보시다시피 같은 질문임에도 똑같은 노드 순환을 반복하며 시간과 토큰이 소모되는 모습입니다.
2. InMemoryCache (/w Langchain)
langchain에서 캐싱을 위한 라이브러리를 제공해주어서 사용해보았습니다.
from langchain.globals import set_llm_cache
from langchain_community.cache import InMemoryCache
class GraphBuilder:
def __init__(self):
set_llm_cache(InMemoryCache())
self.graph = self._create_graph()
self.tracer = OpikTracer(
graph=self.graph.get_graph(xray=True),
project_name="agiledocs-rag"
)
위와같은 코드를 사용해 graph를 생성하기전에, InMemoryCache를 생성해주었습니다. 적용전 이미지는 하단과 같습니다.
set_llm_cache(InMemoryCache()) 적용 후의 이미지입니다.
보시면 query_classifier node에서는 0초의 시간이 소모되지만, 다른 노드는 여전히 순환하는 모습을 볼 수 있습니다. 이는 Langchain에서 제공하는 set_llm_cache(InMemoryCache())가 LLM 호출만 캐싱을 수행하기 때문입니다. 그렇다면 검색과 리랭킹을 제외한 다른 부분에서는 캐싱이 이루어져야하지않나? 라는 생각이 들 수 있지만, 멀티턴 대화 구성을 위해 프롬프트에 이전 메시지들인 chat_history가 새롭게 누적되다보니 동적인 프롬프트가 되어서 캐싱이 이루어지지 않는 것 입니다.
3. InMemoryCache (/w LangGraph)
그렇다면 Langchain을 활용한 캐싱이 최선일까? 라는 생각에서 시작해서 LLM을 통해 생성한 답변은 캐싱하지 못하더라도 retrieve, rerank처럼 동일한 역할을 수행하는 정적인 노드에 대해서만 이라도 캐싱을 수행할 수는 없을까? 라는 생각이 들어서 찾아낸 것이 LangGraph에서 새롭게 고안해낸 노드 단위의 캐싱입니다. 아래는 샘플 코드입니다.
from langgraph.cache.memory import InMemoryCache
from langgraph.types import CachePolicy
import hashlib
class Cache:
def query_classifier_cache_key(self,state : State) -> str:
"""쿼리 분류기용 캐시 키 - 쿼리 내용만 사용"""
query = state.get("query", "")
normalized_query = query.strip().lower()
cache_key = hashlib.md5(normalized_query.encode()).hexdigest()
return cache_key
workflow.add_node("query_classifier", queryclassifier.classification, cache_policy=CachePolicy(key_func=cache.query_classifier_cache_key))
workflow.compile(cache=InMemoryCache())
캐싱을 하고자하는 노드에 CachePolicy()를 사용해주었고, 인자로 ttl과 key_func를 사용할 수 있습니다. ttl은 캐싱 유효시간을 의미하고 key_func는 캐싱의 키값을 반환하는 함수를 호출합니다.
구체적인 사용 흐름은 다음과 같습니다.
- 노드 실행전 LangGraph가 key_func를 호출해서 캐시 키 생성
- 생성된 키로 이전 결과 검색
- 값이 있으면 노드 실행 skip 후 바로 결과 반환, 없으면 노드 실행후 결과를 캐시에 저장
동일한 쿼리에 대해서는 항상 동일한 캐시 키값을 생성하도록 코드를 작성했기 때문에, 이전에 같은 질문을 했다면 캐싱된 값을 반환하게됩니다. 적용전 이미지는 하단과 같습니다.
적용 후의 이미지입니다.
시간도 토큰도 확실히 절약된 모습을 볼 수 있었습니다. 여기서 Langchain의 캐싱과 다른 모습을 발견할 수 있는데, LangGraph의 캐싱은 일치하는 키값이 있으면 바로 값을 반환하기때문에 노드를 아에 순환도 하지 않습니다. 반면, Langchain의 캐싱은 0초이긴하지만 노드 자체는 순환한다는 점에서 차별점이 있습니다.
그렇다면 노드를 아에 순환도 하지않는데 캐싱으로 반환하는 값이 무엇이길래 응답을 잘 하는 것이지?하고 생각이 들어서 이 부분도 확인해보았습니다. 처음에는 Multi-turn으로 구성하다보니 이전 대화에서 사용한 state가 다음 대화 기록에 사용 돼서 응답이 정상적으로 이루어지는건 아닐까? 하고 의심했었습니다. 그래서 질문 A,B가 있다고할때 A->B->A 순서로 질문을 해서 B의 state를 담은 상태에서 다시 A질문을 해보았습니다. 그런데도 정상 응답이 나오는걸 확인하고, 이전에 어떤 질문을 하던 상관없이 동일한 질문에 해당하는 값을 캐싱으로 반환할때 이전에 생성했던 상태값을 그대로 반환해준다는 것을 알게되었습니다.
마지막으로 아래는 Langchain, Langgraph의 캐싱을 모두 적용시킨 후 캐싱이 됐을 때 모습입니다.
query_classifier node가 나오지 않는 모습을 보아, LangGraph의 캐싱이 Langchain의 캐싱보다 우선수위가 높다는 것을 알 수 있었습니다.
참고자료
https://python.langchain.com/docs/how_to/llm_caching/